Skip to main content
`Ctrl`+`K`

# Source code for astropy.uncertainty.core

    # Licensed under a 3-clause BSD style license - see LICENSE.rst

    """
    Distribution class and associated machinery.
    """

    import builtins

    import numpy as np

    from astropy import stats
    from astropy import units as u
    from astropy.utils.compat.numpycompat import NUMPY_LT_2_0

    if NUMPY_LT_2_0:
        from numpy.core.multiarray import normalize_axis_index
        from numpy.lib.function_base import _parse_gufunc_signature
        from numpy.lib.stride_tricks import DummyArray
    else:
        from numpy.lib._function_base_impl import _parse_gufunc_signature
        from numpy.lib._stride_tricks_impl import DummyArray
        from numpy.lib.array_utils import normalize_axis_index

    from .function_helpers import FUNCTION_HELPERS

    __all__ = ["Distribution"]

    # we set this by hand because the symbolic expression (below) requires scipy
    # SMAD_SCALE_FACTOR = 1 / scipy.stats.norm.ppf(0.75)
    SMAD_SCALE_FACTOR = 1.48260221850560203193936104071326553821563720703125

    class Distribution:
        """A scalar value or array values with associated uncertainty distribution.

        This object will take its exact type from whatever the ``samples``
        argument is. In general this is expected to be ``NdarrayDistribution`` for
        |ndarray| input, and, e.g., ``QuantityDistribution`` for a subclass such
        as |Quantity|. But anything compatible with `numpy.asanyarray` is possible
        (generally producing ``NdarrayDistribution``).

        See also: https://docs.astropy.org/en/stable/uncertainty/

        Parameters
        ----------
        samples : array-like
            The distribution, with sampling along the *trailing* axis. If 1D, the sole
            dimension is used as the sampling axis (i.e., it is a scalar distribution).
            If an |ndarray| or subclass, the data will not be copied unless it is not
            possible to take a view (generally, only when the strides of the last axis
            are negative).

        """

        _generated_subclasses = {}

        def __new__(cls, samples):
            if isinstance(samples, Distribution):
                samples = samples.distribution
            else:
                # We can handle strides larger than the itemsize by stride trickery,
                # but the samples do have to have positive offset from each other,
                # so if that's not the case, we make a copy.
                if not isinstance(samples, np.ndarray) or (
                    samples.strides[-1] < samples.dtype.itemsize
                    and samples.shape[-1:] != (1,)  # strides[-1] can be 0.
                ):
                    samples = np.asanyarray(samples, order="C")
                if samples.shape == ():
                    raise TypeError("Attempted to initialize a Distribution with a scalar")

            # Do the view in two stages, since for viewing as a new class, it is good
            # to have with the dtype in place (e.g., for Quantity.__array_finalize__,
            # it is needed to create the correct StructuredUnit).
            # First, get the structured dtype with a single "samples" field that holds
            # a size n_samples array of "sample" fields with individual samples.
            # We need the double indirection to be able to deal with non-contiguous samples
            # (such as result from indexing a structured array).
            # Note that for a single sample, samples.strides[-1] can be 0.
            new_dtype = cls._get_distribution_dtype(
                samples.dtype, samples.shape[-1], itemsize=samples.strides[-1]
            )
            # Create a structured array with the new dtype. We ensure we start with a
            # regular ndarray, to avoid interference between our complicated dtype and
            # something else (such as would be the case for Quantity.__array_finalize__,
            # which needs the dtype to be correct for creating a StructuredUnit).
            if samples.dtype.itemsize == samples.strides[-1]:
                # For contiguous last axis, a plain new view suffices.
                structured = samples.view(np.ndarray)
            else:
                # But if the last axis is not contiguous, we use __array_interface__ (as in
                # np.lib.stridetricks.as_strided) to create a new ndarray where the last
                # axis is covered by a void of the same size as the structured new_dtype
                # (assigning it directly does not work: possible empty fields get entries).
                interface = dict(samples.__array_interface__)
                interface["descr"] = interface["typestr"] = f"|V{new_dtype.itemsize}"
                interface["shape"] = samples.shape[:-1]
                interface["strides"] = samples.strides[:-1]
                structured = np.asarray(DummyArray(interface, base=samples))
            # Set our new structured dtype.
            structured.dtype = new_dtype
            # Get rid of trailing dimension of 1.
            structured.shape = samples.shape[:-1]

            # Now view as the Distribution subclass, and finalize based on the
            # original samples (e.g., to set the unit for QuantityDistribution).
            new_cls = cls._get_distribution_cls(type(samples))
            self = structured.view(new_cls)
            self.__array_finalize__(samples)
            return self

        @classmethod
        def _get_distribution_cls(cls, samples_cls):
            if issubclass(samples_cls, Distribution):
                return samples_cls

            new_cls = cls._generated_subclasses.get(samples_cls)
            if new_cls is None:
                # Walk through MRO and find closest base data class.
                # Note: right now, will basically always be ndarray, but
                # one could imagine needing some special care for one subclass,
                # which would then get its own entry.  E.g., if MaskedAngle
                # defined something special, then MaskedLongitude should depend
                # on it.
                for mro_item in samples_cls.__mro__:
                    base_cls = cls._generated_subclasses.get(mro_item)
                    if base_cls is not None:
                        break
                else:
                    # Just hope that NdarrayDistribution can handle it.
                    # TODO: this covers the case where a user puts in a list or so,
                    # but for those one could just explicitly do something like
                    # _generated_subclasses[list] = NdarrayDistribution
                    return NdarrayDistribution

                # We need to get rid of the top _DistributionRepr, since we add
                # it again below (it should always be on top).
                # TODO: remove the need for _DistributionRepr by defining
                # __array_function__ and overriding array2string.
                if base_cls.__mro__[1] is not _DistributionRepr:
                    # Sanity check. This should never happen!
                    raise RuntimeError(
                        f"found {base_cls=}, which does not have _DistributionRepr at "
                        "__mro__[1]. Please raise an issue at "
                        "https://github.com/astropy/astropy/issues/new/choose."
                    )
                # Create (and therefore register) new Distribution subclass for the
                # given samples_cls.
                new_cls = type(
                    samples_cls.__name__ + "Distribution",
                    (_DistributionRepr, samples_cls) + base_cls.__mro__[2:],
                    {"_samples_cls": samples_cls},
                )
                cls._generated_subclasses[samples_cls] = new_cls

            return new_cls

        @classmethod
        def _get_distribution_dtype(cls, dtype, n_samples, itemsize=None):
            dtype = np.dtype(dtype)
            # If not a sample dtype already, create one with a "samples" entry that holds
            # all the samples.  Here, we create an indirection via "sample" because
            # selecting an item from a structured array will necessarily give
            # non-consecutive samples, and these can only be dealt with samples that are
            # larger than a single element (plus stride tricks; see __new__). For those
            # cases, itemsize larger than dtype.itemsize will be passed in (note that for
            # the n_sample=1 case, itemsize can be 0; like itemsize=None, this will be dealt
            # with by "or dtype.itemsize" below).
            if dtype.names != ("samples",):
                sample_dtype = np.dtype(
                    dict(
                        names=["sample"],
                        formats=[dtype],
                        itemsize=itemsize or dtype.itemsize,
                    )
                )
                dtype = np.dtype([("samples", sample_dtype, (n_samples,))])
            return dtype

        @property
        def distribution(self):
            return self["samples"]["sample"]

        @property
        def dtype(self):
            return super().dtype["samples"].base["sample"]

        @dtype.setter
        def dtype(self, dtype):
            dtype = self._get_distribution_dtype(
                dtype, self.n_samples, itemsize=super().dtype["samples"].base.itemsize
            )
            super(Distribution, self.__class__).dtype.__set__(self, dtype)

        def astype(self, dtype, *args, **kwargs):
            dtype = self._get_distribution_dtype(dtype, self.n_samples)
            return super().astype(dtype, *args, **kwargs)

        def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
            converted = []
            outputs = kwargs.pop("out", None)
            if outputs:
                kwargs["out"] = tuple(
                    (output.distribution if isinstance(output, Distribution) else output)
                    for output in outputs
                )
                if ufunc.nout == 1:
                    outputs = outputs[0]

            axis = kwargs.get("axis")
            keepdims = kwargs.get("keepdims", False)
            if method in {"reduce", "accumulate", "reduceat"}:
                if axis is None:
                    assert isinstance(inputs[0], Distribution)
                    kwargs["axis"] = tuple(range(inputs[0].ndim))

            for input_ in inputs:
                # For any input that is not a Distribution, we add an axis at the
                # end, to allow proper broadcasting with the distributions.
                if isinstance(input_, Distribution):
                    converted.append(input_.distribution)
                else:
                    shape = getattr(input_, "shape", ())
                    if shape:
                        converted.append(input_[..., np.newaxis])
                    else:
                        converted.append(input_)

            if ufunc.signature:
                # We're dealing with a gufunc, which may add an extra axis.
                # Ignore axes keyword for now...  TODO: remove this limitation.
                if "axes" in kwargs:
                    raise NotImplementedError(
                        "Distribution does not yet support gufunc calls with 'axes'."
                    )
                # Parse signature with private numpy function. Note it
                # cannot handle spaces in tuples, so remove those.
                in_sig, out_sig = _parse_gufunc_signature(ufunc.signature.replace(" ", ""))
                ncore_in = [len(sig) for sig in in_sig]
                ncore_out = [len(sig) + (1 if keepdims else 0) for sig in out_sig]
                if ufunc.nout == 1:
                    ncore_out = ncore_out[0]
                if axis is None:
                    axis = -1
                converted = [
                    np.moveaxis(conv, -1, normalize_axis_index(axis, conv.ndim) - ncore)
                    if ncore and getattr(conv, "shape", ())
                    else conv
                    for conv, ncore in zip(converted, ncore_in)
                ]
            else:
                ncore_out = None

            results = getattr(ufunc, method)(*converted, **kwargs)

            return self._result_as_distribution(results, outputs, ncore_out, axis)

        def __array_function__(self, function, types, args, kwargs):
            # TODO: go through functions systematically to see which ones
            # work and/or can be supported.
            if function in FUNCTION_HELPERS:
                function_helper = FUNCTION_HELPERS[function]
                try:
                    args, kwargs, out = function_helper(*args, **kwargs)
                except NotImplementedError:
                    return self._not_implemented_or_raise(function, types)

                result = super().__array_function__(function, types, args, kwargs)
                # Fall through to return section, wrapping distributions, unless
                # indicated otherwise.  TODO: use a less hacky solution?
                if out is True:
                    return result

            else:  # pragma: no cover
                # By default, just pass it through for now.
                return super().__array_function__(function, types, args, kwargs)

            # We're done if the result was NotImplemented, which can happen
            # if other inputs/outputs override __array_function__;
            # hopefully, they can then deal with us.
            if result is NotImplemented:
                return NotImplemented

            return self._result_as_distribution(result, out=out)

        def _result_as_distribution(self, result, out, ncore_out=None, axis=None):
            """Turn result into a distribution.

            If no output is given, it will create a Distribution from the array,
            If an output is given, it should be fine as is.

            Parameters
            ----------
            result : ndarray or tuple thereof
                Array(s) which need to be turned into Distribution.
            out : Distribution, tuple of Distribution or None
                Possible output |Distribution|. Should be `None` or a tuple if result
                is a tuple.
            ncore_out: int or tuple thereof
                The number of core dimensions for the output array for a gufunc.  This
                is used to determine which axis should be used for the samples.
            axis: int or None
                The axis a gufunc operated on.  Used only if ``ncore_out`` is given.

            Returns
            -------
            out : Distribution
            """
            if isinstance(result, (tuple, list)):
                if out is None:
                    out = (None,) * len(result)
                if ncore_out is None:
                    ncore_out = (None,) * len(result)
                # Some np.linalg functions return namedtuple, which is handy to access
                # elements by name, but cannot be directly initialized with an iterator.
                result_cls = getattr(result, "_make", result.__class__)
                return result_cls(
                    self._result_as_distribution(result_, out_, axis=axis, ncore_out=ncore)
                    for (result_, out_, ncore) in zip(result, out, ncore_out)
                )

            if out is None:
                # Turn the result into a Distribution if needed.
                if not isinstance(result, Distribution) and getattr(result, "shape", ()):
                    if ncore_out is not None:
                        result = np.moveaxis(
                            result, normalize_axis_index(axis, result.ndim) - ncore_out, -1
                        )
                    return Distribution(result)
                else:
                    return result
            else:
                # TODO: remove this sanity check once test cases are more complete.
                assert isinstance(out, Distribution)
                return out

        def _not_implemented_or_raise(self, function, types):
            # Our function helper or dispatcher found that the function does not
            # work with Distribution.  In principle, there may be another class that
            # knows what to do with us, for which we should return NotImplemented.
            # But if there is ndarray (or a non-Distribution subclass of it) around,
            # it quite likely coerces, so we should just break.
            if any(
                issubclass(t, np.ndarray) and not issubclass(t, Distribution) for t in types
            ):
                raise TypeError(
                    f"the Distribution implementation cannot handle {function} "
                    "with the given arguments."
                ) from None
            else:
                return NotImplemented

        # Override __eq__ and __ne__ to pass on directly to the ufunc since
        # otherwise comparisons with non-distributions do not work (but
        # deferring if other defines __array_ufunc__ = None -- see
        # numpy/core/src/common/binop_override.h for the logic; we assume we
        # will never deal with __array_priority__ any more).  Note: there is no
        # problem for other comparisons, since for those, structured arrays are
        # not treated differently in numpy/core/src/multiarray/arrayobject.c.
        def __eq__(self, other):
            if getattr(other, "__array_ufunc__", False) is None:
                return NotImplemented
            return np.equal(self, other)

        def __ne__(self, other):
            if getattr(other, "__array_ufunc__", False) is None:
                return NotImplemented
            return np.not_equal(self, other)

        @property
        def n_samples(self):
            """
            The number of samples of this distribution.  A single `int`.
            """
            return super().dtype["samples"].shape[0]

        def pdf_mean(self, dtype=None, out=None):
            """
            The mean of this distribution.

            Arguments are as for `numpy.mean`.
            """
            return self.distribution.mean(axis=-1, dtype=dtype, out=out)

        def pdf_std(self, dtype=None, out=None, ddof=0):
            """
            The standard deviation of this distribution.

            Arguments are as for `numpy.std`.
            """
            return self.distribution.std(axis=-1, dtype=dtype, out=out, ddof=ddof)

        def pdf_var(self, dtype=None, out=None, ddof=0):
            """
            The variance of this distribution.

            Arguments are as for `numpy.var`.
            """
            return self.distribution.var(axis=-1, dtype=dtype, out=out, ddof=ddof)

        def pdf_median(self, out=None):
            """
            The median of this distribution.

            Parameters
            ----------
            out : array, optional
                Alternative output array in which to place the result. It must
                have the same shape and buffer length as the expected output,
                but the type (of the output) will be cast if necessary.
            """
            return np.median(self.distribution, axis=-1, out=out)

        def pdf_mad(self, out=None):
            """
            The median absolute deviation of this distribution.

            Parameters
            ----------
            out : array, optional
                Alternative output array in which to place the result. It must
                have the same shape and buffer length as the expected output,
                but the type (of the output) will be cast if necessary.
            """
            median = self.pdf_median(out=out)
            absdiff = np.abs(self - median)
            return np.median(
                absdiff.distribution, axis=-1, out=median, overwrite_input=True
            )

        def pdf_smad(self, out=None):
            """
            The median absolute deviation of this distribution rescaled to match the
            standard deviation for a normal distribution.

            Parameters
            ----------
            out : array, optional
                Alternative output array in which to place the result. It must
                have the same shape and buffer length as the expected output,
                but the type (of the output) will be cast if necessary.
            """
            result = self.pdf_mad(out=out)
            result *= SMAD_SCALE_FACTOR
            return result

        def pdf_percentiles(self, percentile, **kwargs):
            """
            Compute percentiles of this Distribution.

            Parameters
            ----------
            percentile : float or array of float or `~astropy.units.Quantity`
                The desired percentiles of the distribution (i.e., on [0,100]).
                `~astropy.units.Quantity` will be converted to percent, meaning
                that a ``dimensionless_unscaled`` `~astropy.units.Quantity` will
                be interpreted as a quantile.

            Additional keywords are passed into `numpy.percentile`.

            Returns
            -------
            percentiles : `~astropy.units.Quantity` ['dimensionless']
                The ``fracs`` percentiles of this distribution.
            """
            percentile = u.Quantity(percentile, u.percent).value
            percs = np.percentile(self.distribution, percentile, axis=-1, **kwargs)
            # numpy.percentile strips units for unclear reasons, so we have to make
            # a new object with units
            if hasattr(self.distribution, "_new_view"):
                return self.distribution._new_view(percs)
            else:
                return percs

        def pdf_histogram(self, **kwargs):
            """
            Compute histogram over the samples in the distribution.

            Parameters
            ----------
            All keyword arguments are passed into `astropy.stats.histogram`. Note
            That some of these options may not be valid for some multidimensional
            distributions.

            Returns
            -------
            hist : array
                The values of the histogram. Trailing dimension is the histogram
                dimension.
            bin_edges : array of dtype float
                Return the bin edges ``(length(hist)+1)``. Trailing dimension is the
                bin histogram dimension.
            """
            distr = self.distribution
            raveled_distr = distr.reshape(distr.size // distr.shape[-1], distr.shape[-1])

            nhists = []
            bin_edges = []
            for d in raveled_distr:
                nhist, bin_edge = stats.histogram(d, **kwargs)
                nhists.append(nhist)
                bin_edges.append(bin_edge)

            nhists = np.array(nhists)
            nh_shape = self.shape + (nhists.size // self.size,)
            bin_edges = np.array(bin_edges)
            be_shape = self.shape + (bin_edges.size // self.size,)
            return nhists.reshape(nh_shape), bin_edges.reshape(be_shape)

    class ScalarDistribution(Distribution, np.void):
        """Scalar distribution.

        This class mostly exists to make `~numpy.array2print` possible for
        all subclasses.  It is a scalar element, still with n_samples samples.
        """

    class ArrayDistribution(Distribution, np.ndarray):
        # This includes the important override of view and __getitem__
        # which are needed for all ndarray subclass Distributions, but not
        # for the scalar one.
        _samples_cls = np.ndarray

        # Override view so that we stay a Distribution version of the new type.
        def view(self, dtype=None, type=None):
            """New view of array with the same data.

            Like `~numpy.ndarray.view` except that the result will always be a new
            `~astropy.uncertainty.Distribution` instance.  If the requested
            ``type`` is a `~astropy.uncertainty.Distribution`, then no change in
            ``dtype`` is allowed.

            """
            if type is None:
                if isinstance(dtype, builtins.type) and issubclass(dtype, np.ndarray):
                    type = self._get_distribution_cls(dtype)
                    dtype = None
                else:
                    type = self.__class__
            else:
                type = self._get_distribution_cls(type)

            type = self._get_distribution_cls(type)
            if dtype is None:
                return super().view(type=type)

            dtype = np.dtype(dtype)
            if dtype == self.dtype:
                return super().view(type=type)

            if dtype.names == ("samples",):
                # Assume the user knows what they are doing.
                return super().view(dtype, type)

            if dtype.shape == () and dtype.itemsize == self.dtype.itemsize:
                dtype = self._get_distribution_dtype(
                    dtype,
                    self.n_samples,
                    itemsize=super(Distribution, self).dtype["samples"].base.itemsize,
                )
                return super().view(dtype, type)

            samples_cls = type._samples_cls
            if dtype.itemsize == self.dtype.itemsize:
                distr = self.distribution
                distr_view = distr.view(dtype, samples_cls)
                # This does not necessarily leave the sample axis in the right place.
                return Distribution(np.moveaxis(distr_view, distr.ndim - 1, -1))
            elif dtype.itemsize == super(Distribution, self).dtype["samples"].base.itemsize:
                distr = np.moveaxis(self.distribution, -1, -2)
                distr_view = distr.view(dtype, samples_cls).squeeze(-1)
                return Distribution(distr_view)
            else:
                raise ValueError(
                    f"{self.__class__} can only be viewed with a dtype with "
                    "itemsize {self.strides[-1]} or {self.dtype.itemsize}"
                )

        @property
        def distribution(self):
            # Like in the creation, we go through an ndarray to ensure we have our
            # actual dtype and to avoid entering, e.g., Quantity.__getitem__, which
            # would give problems with units.
            structured = super().view(np.ndarray)
            distribution = structured["samples"]["sample"].view(self._samples_cls)
            distribution.__array_finalize__(self)
            return distribution

        def __getitem__(self, item):
            if isinstance(item, str):
                # "samples" should always get back to the samples class.
                if item == "samples":
                    return self.distribution
                else:
                    # Hard to get this right directly, so instead get item from the
                    # distribution, and create a new instance.  We move the sample axis to
                    # the end to ensure the order is right for possible subarrays.
                    return Distribution(np.moveaxis(self.distribution[item], self.ndim, -1))

            if isinstance(item, Distribution):
                # Required for in-place operations like dist[dist < 0] += 360.
                return self.distribution[item.distribution]

            result = super().__getitem__(item)
            if isinstance(result, np.void):
                return result.view((ScalarDistribution, result.dtype))
            else:
                return result

        def __setitem__(self, item, value):
            if isinstance(item, Distribution):
                # Support operations like dist[dist < 0] = 0.
                self.distribution[item.distribution] = value
                return

            if isinstance(item, str):
                if item == "samples":
                    self.distribution[()] = value
                    return

                # Get a view of this item (non-trivial; see above).
                self = self[item]
                item = ()

            if not isinstance(value, Distribution):
                # If value is not already a Distribution, first make it an array
                # to help interpret possible structured dtype, and then turn it
                # into a Distribution with n_samples=1 (which will broadcast).
                value = np.asanyarray(value, dtype=self.dtype)
                value = Distribution(value[..., np.newaxis])

            super().__setitem__(item, value)

    class _DistributionRepr:
        def __repr__(self):
            reprarr = repr(self.distribution)
            if reprarr.endswith(">"):
                firstspace = reprarr.find(" ")
                reprarr = reprarr[firstspace + 1 : -1]  # :-1] removes the ending '>'
                return (
                    f"<{self.__class__.__name__} {reprarr} with n_samples={self.n_samples}>"
                )
            else:  # numpy array-like
                firstparen = reprarr.find("(")
                reprarr = reprarr[firstparen:]
                return f"{self.__class__.__name__}{reprarr} with n_samples={self.n_samples}"

        def __str__(self):
            distrstr = str(self.distribution)
            toadd = f" with n_samples={self.n_samples}"
            return distrstr + toadd

        def _repr_latex_(self):
            if hasattr(self.distribution, "_repr_latex_"):
                superlatex = self.distribution._repr_latex_()
                toadd = rf", \; n_{{\rm samp}}={self.n_samples}"
                return superlatex[:-1] + toadd + superlatex[-1]
            else:
                return None

    class NdarrayDistribution(_DistributionRepr, ArrayDistribution):
        pass

    # Ensure our base NdarrayDistribution is known.
    Distribution._generated_subclasses[np.ndarray] = NdarrayDistribution