gwkokab.utils.math
==================

.. py:module:: gwkokab.utils.math


Functions
---------

.. autoapisummary::

   gwkokab.utils.math.beta_dist_concentrations_to_mean_variance
   gwkokab.utils.math.beta_dist_mean_variance_to_concentrations
   gwkokab.utils.math.cumtrapz
   gwkokab.utils.math.logsubexp
   gwkokab.utils.math.truncnorm_logpdf


Module Contents
---------------

.. py:function:: beta_dist_concentrations_to_mean_variance(alpha: jaxtyping.ArrayLike, beta: jaxtyping.ArrayLike, loc: jaxtyping.ArrayLike = 0.0, scale: jaxtyping.ArrayLike = 1.0) -> Tuple[jaxtyping.ArrayLike, jaxtyping.ArrayLike]

   Let :math:`\alpha` and :math:`\beta` be the shape parameters of a beta
   distribution, :math:`a` being the location and :math:`b` being the scale. This
   function returns the mean and variance of the distribution. Then concentrations are
   given by:

   .. math::
       \mu = a+b\frac{\alpha}{\alpha + \beta}\qquad
       \sigma^2 = b^2\frac{\alpha \beta}{(\alpha + \beta)^2 (\alpha + \beta + 1)}

   :param alpha: The shape parameter :math:`\alpha`.
   :type alpha: ArrayLike
   :param beta: The shape parameter :math:`\beta`.
   :type beta: ArrayLike
   :param loc: The location :math:`a` of the beta distribution.
   :type loc: ArrayLike
   :param scale: The scale :math:`b` of the beta distribution.
   :type scale: ArrayLike

   :returns: The mean :math:`\mu` and variance :math:`\sigma^2` of the beta distribution.
   :rtype: Tuple[ArrayLike, ArrayLike]


.. py:function:: beta_dist_mean_variance_to_concentrations(mean: jaxtyping.ArrayLike, variance: jaxtyping.ArrayLike, loc: jaxtyping.ArrayLike = 0.0, scale: jaxtyping.ArrayLike = 1.0) -> Tuple[jaxtyping.ArrayLike, jaxtyping.ArrayLike]

   Let :math:`\mu` and :math:`\sigma^2` be the mean and variance of a beta
   distribution, :math:`a` being the location and :math:`b` being the scale. This
   function returns the shape parameters :math:`\alpha` and :math:`\beta` of the
   distribution. Then concentrations are given by:

   .. math::
       \alpha = -\frac{\mu-a}{b} \left(\left(\frac{\mu-a}{\sigma}\right)\left(\frac{\mu-a-b}{\sigma}\right)+1\right)\qquad
       \beta = \alpha\left(\frac{b}{\mu-a}-1\right)

   :param mean: The mean :math:`\mu` of the beta distribution.
   :type mean: ArrayLike
   :param variance: The variance :math:`\sigma^2` of the beta distribution.
   :type variance: ArrayLike
   :param loc: The location :math:`a` of the beta distribution.
   :type loc: ArrayLike
   :param scale: The scale :math:`b` of the beta distribution.
   :type scale: ArrayLike

   :returns: The shape parameters :math:`\alpha` and :math:`\beta` of the beta distribution.
   :rtype: Tuple[ArrayLike, ArrayLike]


.. py:function:: cumtrapz(y: jaxtyping.Array, x: jaxtyping.Array) -> jaxtyping.Array

   Calculate the cumulative trapezoidal integration of y with respect to x.

   :param y: array to integrate
   :type y: Array
   :param x: array to integrate over
   :type x: Array

   :returns: The result of the cumulative trapezoidal integration of y with respect to x.
   :rtype: Array


.. py:function:: logsubexp(a: jaxtyping.ArrayLike, b: jaxtyping.ArrayLike) -> jaxtyping.ArrayLike

   Compute :math:`\log(\exp(a) - \exp(b))` in a numerically stable way.

   :param a: input array
   :type a: ArrayLike
   :param b: input array
   :type b: ArrayLike

   :returns: The value of :math:`\log(\exp(a) - \exp(b))`.
   :rtype: ArrayLike


.. py:function:: truncnorm_logpdf(xx: jaxtyping.ArrayLike, loc: jaxtyping.ArrayLike, scale: jaxtyping.ArrayLike, low: jaxtyping.ArrayLike, high: jaxtyping.ArrayLike) -> jaxtyping.ArrayLike

   Compute the log probability density function of a truncated normal distribution.

   :param xx: The input array.
   :type xx: ArrayLike
   :param loc: The mean of the distribution.
   :type loc: ArrayLike
   :param scale: The standard deviation of the distribution.
   :type scale: ArrayLike
   :param low: The lower bound of the distribution.
   :type low: ArrayLike
   :param high: The upper bound of the distribution.
   :type high: ArrayLike

   :returns: The log probability density function of the truncated normal distribution.
   :rtype: ArrayLike