jexplore.tools.distributions ============================ .. py:module:: jexplore.tools.distributions .. autoapi-nested-parse:: Distributions ============= Classes for sampling and evaluating some relevant distributions. Attributes ---------- .. autoapisummary:: jexplore.tools.distributions.DrawFn Classes ------- .. autoapisummary:: jexplore.tools.distributions.Distr jexplore.tools.distributions.Uniform jexplore.tools.distributions.Normal jexplore.tools.distributions.MVNormal jexplore.tools.distributions.StudentT Module Contents --------------- .. py:type:: DrawFn :canonical: Callable[[jax.Array, tuple], tuple[jax.Array, jax.Array]] Prototype definition of a drawing function. .. py:class:: Distr(dim) Bases: :py:obj:`Protocol` Abstract parent distribution class. :param int dim: space dimension .. py:attribute:: dim :type: int .. py:method:: sample(key, shape) Samples the distribution. :param ArrayLike key: PRNG key used as the random key. :param tuple shape: shape of the sample. :return: the actualized PRNG key and the samples with shape shape + (py:attr:`dim`,) :rtype: tuple .. py:method:: eval(x) Evaluates the distribution on a set of points :param array x: points array, with shape (..., py:attr:`dim`) :return: the values of the distribution on the given points :rtype: Array .. py:method:: leval(x) Evaluates the log distribution on a set of points :param array x: points array, with shape (..., py:attr:`dim`) :return: the values of the log distribution on the given points :rtype: Array .. py:class:: Uniform(dim, minval = 0.0, maxval = 1.0) Bases: :py:obj:`Distr` Constant distribution in a box: :param int dim: space dimension :param RealArray minval: minimum (inclusive) value broadcast-compatible with shape for the range (default 0). :param RealArray maxval: maximum (exclusive) value broadcast-compatible with shape for the range (default 1). .. py:attribute:: minval :type: jax.Array .. py:attribute:: maxval :type: jax.Array .. py:attribute:: lvol :type: jax.Array .. py:method:: inbox(x) Check if values are in the support of the distribution :param x: values to be checked (last dimension of the shape should be dimension of the space on which the distribution is defined. :return: the boolean array of the results of the test. .. py:method:: sample(key, shape) Samples the distribution. :param ArrayLike key: PRNG key used as the random key. :param tuple shape: shape of the sample. :return: the actualized PRNG key and the samples with shape shape + (py:attr:`dim`,) :rtype: tuple .. py:method:: leval(x) Evaluates the log distribution on a set of points :param array x: points array, with shape (..., py:attr:`dim`) :return: the values of the log distribution on the given points :rtype: Array .. py:method:: eval(x) Evaluates the distribution on a set of points :param array x: points array, with shape (..., py:attr:`dim`) :return: the values of the distribution on the given points :rtype: Array .. py:class:: Normal(dim) Bases: :py:obj:`Distr` Normal distribution with identity covariance. :param int dim: space dimension .. py:attribute:: lnorm :type: jax.Array .. py:method:: sample(key, shape) Samples the distribution. :param ArrayLike key: PRNG key used as the random key. :param tuple shape: shape of the sample. :return: the actualized PRNG key and the samples with shape shape + (py:attr:`dim`,) :rtype: tuple .. py:method:: leval(x) Evaluates the log distribution on a set of points :param array x: points array, with shape (..., py:attr:`dim`) :return: the values of the log distribution on the given points :rtype: Array .. py:class:: MVNormal(cov, scalar = False, norm=True) Bases: :py:obj:`Normal` Multivariate Normal distribution :param cov: covariance matrix :param scalar: if true the distribution evaluation is not vectorized (this is useful to use the distribution as a likelihood). :param norm: if true the evaluation returns normalized values. .. py:attribute:: cov :type: jax.Array Set of covariance matrices .. py:attribute:: icov :type: jax.Array Inverse covariance matrices .. py:attribute:: lower :type: jax.Array Cholesky lower matrices .. py:attribute:: scalar :type: bool .. py:attribute:: lnorm .. py:method:: sample(key, shape) Samples the distribution. :param ArrayLike key: PRNG key used as the random key. :param tuple shape: shape of the sample. :return: the actualized PRNG key and the samples with shape shape + (py:attr:`dim`,) :rtype: tuple .. py:method:: leval(x) Evaluates the log distribution on a set of points :param array x: points array, with shape (..., py:attr:`dim`) :return: the values of the log distribution on the given points :rtype: Array .. py:method:: draw_cov(key, shape, minev = 0.5, maxev = 1.0) :staticmethod: Draw a random set of positive defined covariance matrix. :param key: PRNG key :param shape: shape of the set of covariance. The last dimension of the tuple must be the dimension of the space. So that each covariance matrix will have shape (shape[-1], shape[-1]) :param minev: lower bound(s) for the eigenvalues :param maxenv: higher bound(s) for the eigenvalues :return: PRNG key and the set of covariance matrices. .. py:class:: StudentT(dim, nu = 5.0) Bases: :py:obj:`Distr` Multivarate student-t with identity covariance. :param int dim: space dimension :param float nu: student-t :math:`\nu` parameter .. py:attribute:: nu :type: float :value: 5.0 .. py:attribute:: normal .. py:method:: sample(key, shape) Samples the distribution. :param ArrayLike key: PRNG key used as the random key. :param tuple shape: shape of the sample. :return: the actualized PRNG key and the samples with shape shape + (py:attr:`dim`,) :rtype: tuple .. py:method:: leval(x) Evaluates the log distribution on a set of points :param array x: points array, with shape (..., py:attr:`dim`) :return: the values of the log distribution on the given points :rtype: Array