Distributions¶

class rv_generic(ABC):
    """Generic random variable class for common functionality."""

    name: ClassVar[Optional[str]] = None
    parameters: ClassVar[Mapping[str, stats_constraints.Constraint]] = {}
    parameter_aliases: ClassVar[Mapping[str, str]] = {}
    extra_frozen_kwds: ClassVar[frozenset[str]] = frozenset()

    def __init__(self, name: Optional[str] = None):
        if name is not None:
            type(self).name = name

    def __call__(self, *args: Any, **kwds: Any) -> "rv_frozen":
        """Call the distribution with the given arguments."""
        return self.freeze(*args, **kwds)

    @staticmethod
    def _format_frozen_class_name(name: Optional[str], fallback: str) -> str:
        """Create a stable class name for generated frozen distributions."""
        raw_name = name or fallback
        parts = "".join(ch if ch.isalnum() else "_" for ch in str(raw_name)).split("_")
        camel = "".join(part.capitalize() for part in parts if part)
        return f"Frozen{camel or 'Distribution'}"

    def _get_or_create_frozen_class(
        self, base_frozen_cls: type["rv_frozen"]
    ) -> type["rv_frozen"]:
        """Get a cached frozen class for this distribution instance."""
        cache = getattr(self, "_frozen_cls_cache", None)
        if cache is not None and issubclass(cache, base_frozen_cls):
            return cast(type["rv_frozen"], cache)

        class_name = self._format_frozen_class_name(
            getattr(self, "name", None), self.__class__.__name__
        )
        frozen_cls = cast(
            type["rv_frozen"],
            FrozenDistributionMeta(
                class_name,
                (base_frozen_cls,),
                {},
                dist=self,
            ),
        )
        self._frozen_cls_cache = frozen_cls
        return frozen_cls

    @staticmethod
    def _distribution_display_name(dist: Any) -> str:
        return (
            getattr(dist, "name", None)
            or getattr(dist, "__name__", None)
            or dist.__class__.__name__
        )

    @classmethod
    def _bind_frozen_args_for_dist(
        cls, dist: Any, args: tuple[Any, ...], kwds: Mapping[str, Any]
    ) -> _FrozenArgs:
        """Validate and bind frozen distribution arguments."""
        parameters = tuple(getattr(dist, "parameters", {}).keys())
        extra_kwds = set(getattr(dist, "extra_frozen_kwds", ()))
        kwds_dict = dict(kwds)
        dist_name = cls._distribution_display_name(dist)

        if len(args) > len(parameters):
            raise TypeError(
                f"{dist_name} expected at most {len(parameters)} positional "
                f"arguments, got {len(args)}."
            )

        positional_names = set(parameters[: len(args)])
        duplicate_names = sorted(positional_names.intersection(kwds_dict))
        if duplicate_names:
            names = ", ".join(repr(name) for name in duplicate_names)
            raise TypeError(f"{dist_name} got multiple values for argument {names}.")

        valid_kwds = set(parameters).union(extra_kwds)
        unexpected = sorted(set(kwds_dict).difference(valid_kwds))
        if unexpected:
            names = ", ".join(repr(name) for name in unexpected)
            raise TypeError(f"{dist_name} got unexpected keyword argument {names}.")

        parameter_values = {
            name: value for name, value in zip(parameters, args, strict=False)
        }
        parameter_values.update({
            name: kwds_dict[name] for name in parameters if name in kwds_dict
        })
        return _FrozenArgs(args=args, kwds=kwds_dict, parameter_values=parameter_values)

    def _freeze_as(
        self, base_frozen_cls: type["rv_frozen"], *args: Any, **kwds: Any
    ) -> "rv_frozen":
        """Freeze a distribution into the requested frozen base class."""
        frozen_cls = self._get_or_create_frozen_class(base_frozen_cls)
        return frozen_cls(self, *args, **kwds)

    def _freeze_as_named(self, cls_name: str, *args: Any, **kwds: Any) -> "rv_frozen":
        """Freeze using a frozen class looked up by name (avoids import cycles)."""
        frozen_cls = cast(type["rv_frozen"], globals()[cls_name])
        return self._freeze_as(frozen_cls, *args, **kwds)

    @classmethod
    def _parse_args(
        cls, *args: Any, **kwds: Any
    ) -> tuple[tuple[Any, ...], dict[str, Any]]:
        """Parse arguments for the distribution."""
        # Move important kwds to args
        params = dict(cls.parameters)
        for param_name in params:
            if param_name in kwds:
                args = args + (kwds.pop(param_name),)
        return args, kwds

    def freeze(self, *args: Any, **kwds: Any) -> "rv_frozen":
        """Freeze the distribution for the given arguments."""
        return self._freeze_as(rv_frozen, *args, **kwds)

    def from_params(
        self, params: Optional[Mapping[str, Any]] = None, **kwds: Any
    ) -> "rv_frozen":
        """Create a frozen distribution from name-keyed parameters."""
        values = dict(params or {})
        values.update(kwds)
        args = []
        for name in self.parameters:
            if name not in values:
                break
            args.append(values.pop(name))
        return self.freeze(*args, **values)

    @classmethod
    def _convert_params(cls, params: Mapping[str, Any], *, to_unconstrained: bool):
        """Shared loop for ``params_to/from_unconstrained`` conversions."""
        converted = {}
        for name, value in params.items():
            constraint = cls.parameters.get(name)
            if isinstance(constraint, stats_constraints.Distribution):
                if to_unconstrained:
                    converted[name] = jax.tree_util.tree_map(
                        _frozen_to_unconstrained,
                        value,
                        is_leaf=_is_frozen_dist,
                    )
                else:
                    converted[name] = jax.tree_util.tree_map(
                        _unconstrained_to_frozen,
                        value,
                        is_leaf=_is_distribution_params,
                    )
                continue
            if not isinstance(constraint, stats_constraints.Constraint):
                converted[name] = value
                continue
            try:
                bijection = biject_to(constraint)
                converted[name] = (
                    bijection.inv(value) if to_unconstrained else bijection(value)
                )
            except NotImplementedError:
                converted[name] = value
        return converted

    @classmethod
    def params_to_unconstrained(cls, params: Mapping[str, Any]) -> dict[str, Any]:
        """Map constrained parameters to an optimization-friendly pytree."""
        return cls._convert_params(params, to_unconstrained=True)

    @classmethod
    def params_from_unconstrained(cls, params: Mapping[str, Any]) -> dict[str, Any]:
        """Map unconstrained parameters back to their declared supports."""
        return cls._convert_params(params, to_unconstrained=False)

    def rvs(
        self,
        rng: RngKey,
        *args: Any,
        shape: Tuple[int, ...] = (),
        name: Optional[str] = None,
        **kwargs: Any,
    ) -> Array:
        """Random variates of given shape.

        Calls through the `rv_p` primitive so traced execution records a random
        variable site while eager execution remains a direct sample.
        """
        return _rv_bind(
            rng,
            args,
            shape,
            self,
            name,
            type(self)._rvs_impl,
            type(self).logpdf,
            kwargs,
        )

    @classmethod
    @abstractmethod
    def support(cls, *args, **kwds) -> stats_constraints.Constraint:
        """Support of the distribution."""
        ...

    @classmethod
    @abstractmethod
    def _rvs_impl(
        cls, rng: RngKey, *args: Any, shape: Tuple[int, ...] = (), **kwargs: Any
    ) -> Array:
        """Implementation for random variate sampling.

        Parameters
        ----------
        rng : jax.random.PRNGKey
            The random key used for sampling
        *args : array_like
            Shape parameters for the distribution
        shape : tuple of ints, optional
            The shape of the samples to draw
        name : str, optional
            Optional site name used when tracing probabilistic programs.
            If omitted, a unique name is generated from the distribution name.
        **kwargs : dict, optional
            Additional parameters (loc, scale, etc.)

        Returns
        -------
        rvs : ndarray or scalar
            Random variates of given shape
        """
        ...

    @classmethod
    def mean(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Mean of the distribution."""
        raise NotImplementedError("Mean is not implemented for this distribution.")

    @classmethod
    def mode(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Mode of the distribution."""
        raise NotImplementedError("Mode is not implemented for this distribution.")

    @classmethod
    def var(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Variance of the distribution."""
        raise NotImplementedError("Variance is not implemented for this distribution.")

    @classmethod
    def std(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Standard deviation of the distribution."""
        return jnp.sqrt(cls.var(*args, **kwds))

    @classmethod
    def cdf(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Cumulative distribution function of the RV."""
        raise NotImplementedError("CDF is not implemented for this distribution.")

    @classmethod
    @abstractmethod
    def logpdf(cls, x: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Log probability of the distribution at the given value."""
        ...

    @classmethod
    def logcdf(cls, x: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Log of the cumulative distribution function at x of the given RV."""
        return jnp.log(cls.cdf(x, *args, **kwds))

    @classmethod
    def sf(cls, x: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Survival function (1 - cdf) at x of the given RV."""
        cdf_val = jnp.asarray(cls.cdf(x, *args, **kwds))
        return jnp.ones_like(cdf_val) - cdf_val

    @classmethod
    def logsf(cls, x: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Log of the survival function at x of the given RV."""
        return jnp.log(cls.sf(x, *args, **kwds))

    @classmethod
    def ppf(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Percent point function (inverse of cdf) of the RV."""
        raise NotImplementedError("PPF is not implemented for this distribution.")

    @classmethod
    def isf(cls, q: ArrayLike, *args: Any, **kwds: Any) -> ArrayLike:
        """Inverse survival function (1 - ppf) of the RV."""
        q = jnp.asarray(q)
        return cls.ppf(1.0 - q, *args, **kwds)

    @classmethod
    def entropy(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Entropy of the RV."""
        raise NotImplementedError("Entropy is not implemented for this distribution.")

    @classmethod
    def median(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Median of the distribution."""
        args, kwds = cls._parse_args(*args, **kwds)
        return cls.ppf(0.5, *args, **kwds)

    @classmethod
    def interval(
        cls, alpha: ArrayLike, *args: Any, **kwds: Any
    ) -> tuple[ArrayLike, ArrayLike]:
        """Confidence interval with equal areas around the median."""
        args, kwds = cls._parse_args(*args, **kwds)
        alpha = jnp.asarray(alpha)
        q1 = (1.0 - alpha) / 2
        q2 = (1.0 + alpha) / 2
        a = cls.ppf(q1, *args, **kwds)
        b = cls.ppf(q2, *args, **kwds)
        return a, b

    @classmethod
    def moment(cls, n: int, *args: Any, **kwds: Any) -> ArrayLike:
        """n-th non-central moment of the distribution.

        Parameters
        ----------
        n : int
            Order of the moment
        *args : array_like
            Shape parameters for the distribution
        **kwds : dict, optional
            Additional parameters (loc, scale, etc.)

        Returns
        -------
        moment : float or ndarray
            n-th non-central moment
        """
        raise NotImplementedError("Moment is not implemented for this distribution.")

    @classmethod
    def skew(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Skewness of the distribution.

        Parameters
        ----------
        *args : array_like
            Shape parameters for the distribution
        **kwds : dict, optional
            Additional parameters (loc, scale, etc.)

        Returns
        -------
        skew : float or ndarray
            Skewness of the distribution
        """
        raise NotImplementedError("Skewness is not implemented for this distribution.")

    @classmethod
    def kurtosis(cls, *args: Any, **kwds: Any) -> ArrayLike:
        """Kurtosis of the distribution.

        Parameters
        ----------
        *args : array_like
            Shape parameters for the distribution
        **kwds : dict, optional
            Additional parameters (loc, scale, etc.)

        Returns
        -------
        kurtosis : float or ndarray
            Kurtosis of the distribution (Fisher's definition, kurtosis - 3)
        """
        raise NotImplementedError("Kurtosis is not implemented for this distribution.")

    @classmethod
    def fit(cls, data: ArrayLike, **kwds: Any) -> tuple[Array, ...]:
        """Maximum likelihood estimation of distribution parameters.

        Parameters
        ----------
        data : array_like
            Data to fit the distribution to
        **kwds : dict, optional
            Additional parameters for the optimization

        Returns
        -------
        params : tuple
            The fitted parameters of the distribution
        """
        raise NotImplementedError("Not implemented for this distribution.")

    @classmethod
    def fit_params(cls, data: ArrayLike, **kwds: Any) -> dict[str, Any]:
        """Fit and return parameters keyed by their declared names."""
        fitted = cls.fit(data, **kwds)
        values = fitted if isinstance(fitted, tuple) else (fitted,)
        return dict(zip(cls.parameters, values, strict=False))

    @classmethod
    def _fit_mle(cls, data: ArrayLike, **kwds: Any) -> tuple[Array, ...]:
        """Fit by optimizing the likelihood in unconstrained parameter space."""
        init_params = {}
        for name in cls.parameters:
            if name not in kwds:
                raise ValueError(f"Provide an initial value for parameter {name!r}.")
            init_params[name] = jnp.asarray(kwds.pop(name))

        initial, unravel = ravel_pytree(cls.params_to_unconstrained(init_params))

        def objective(flat_params):
            params = cls.params_from_unconstrained(unravel(flat_params))
            return -jnp.sum(cls.logpdf(data, **params))

        result = minimize(objective, initial, method="BFGS", **kwds)
        if not bool(jnp.all(jnp.isfinite(result.x))):
            message = getattr(result, "message", "unknown error")
            raise ValueError(f"Optimization failed: {message}")

        fitted = cls.params_from_unconstrained(unravel(result.x))
        return tuple(fitted[name] for name in cls.parameters)

class rv_continuous(rv_generic):
    """Base class for continuous random variables."""

    @classmethod
    def support(cls, *args: Any, **kwds: Any) -> stats_constraints.Constraint:
        """Support of the distribution; the whole real line unless overridden."""
        return stats_constraints.real

    def freeze(self, *args: Any, **kwds: Any) -> "rv_continuous_frozen":
        """Freeze the distribution for the given arguments."""
        return cast(
            "rv_continuous_frozen",
            self._freeze_as(rv_continuous_frozen, *args, **kwds),
        )

    @classmethod
    def pdf(cls, x: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Probability density function at x of the given RV."""
        args, kwds = cls._parse_args(*args, **kwds)
        return jnp.exp(cls.logpdf(x, *args, **kwds))

    @classmethod
    def logpdf(cls, x: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Log of the probability density function at x of the given RV."""
        raise NotImplementedError("Logpdf is not implemented for this distribution.")

    @classmethod
    def fit(cls, data: ArrayLike, **kwds: Any) -> tuple[Array, ...]:
        """Maximum likelihood estimation in unconstrained parameter space."""
        return cls._fit_mle(data, **kwds)

class rv_discrete(rv_generic):
    """Base class for discrete random variables."""

    def freeze(self, *args: Any, **kwds: Any) -> "rv_discrete_frozen":
        """Freeze the distribution for the given arguments."""
        return cast(
            "rv_discrete_frozen",
            self._freeze_as(rv_discrete_frozen, *args, **kwds),
        )

    @classmethod
    def pmf(cls, k: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Probability mass function at k of the given RV."""
        return jnp.exp(cls.logpmf(k, *args, **kwds))

    @classmethod
    def logpmf(cls, k: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Log of the probability mass function at k of the given RV."""
        return jnp.log(cls.pmf(k, *args, **kwds))

    @classmethod
    def logpdf(cls, x: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Log of the probability density function at x of the given RV."""
        return cls.logpmf(x, *args, **kwds)

    @classmethod
    def pdf(cls, x: ArrayLike, *args: Any, **kwds: Any) -> Array:
        """Probability density function at x of the given RV."""
        return jnp.exp(cls.logpdf(x, *args, **kwds))