FM / distributions

/usr/local/lib64/python3.6/site-packages/torch/distributions/

/usr/local/lib64/python3.6/site-packages/torch/distributions

Name	Size	Mode	Actions
__pycache__/	-	0755	rm
bernoulli.py	3904	0644	edit dl rm
beta.py	3406	0644	edit dl rm
binomial.py	5179	0644	edit dl rm
categorical.py	5488	0644	edit dl rm
cauchy.py	2714	0644	edit dl rm
chi2.py	909	0644	edit dl rm
constraints.py	17288	0644	edit dl rm
constraint_registry.py	10234	0644	edit dl rm
continuous_bernoulli.py	8532	0644	edit dl rm
dirichlet.py	3584	0644	edit dl rm
distribution.py	11735	0644	edit dl rm
exponential.py	2525	0644	edit dl rm
exp_family.py	2275	0644	edit dl rm
fishersnedecor.py	3152	0644	edit dl rm
gamma.py	3121	0644	edit dl rm
geometric.py	4266	0644	edit dl rm
gumbel.py	2528	0644	edit dl rm
half_cauchy.py	2257	0644	edit dl rm
half_normal.py	2058	0644	edit dl rm
independent.py	4361	0644	edit dl rm
kl.py	29998	0644	edit dl rm
kumaraswamy.py	2927	0644	edit dl rm
laplace.py	3054	0644	edit dl rm
lkj_cholesky.py	6124	0644	edit dl rm
logistic_normal.py	1983	0644	edit dl rm
log_normal.py	1772	0644	edit dl rm
lowrank_multivariate_normal.py	9930	0644	edit dl rm
mixture_same_family.py	8636	0644	edit dl rm
multinomial.py	4776	0644	edit dl rm
multivariate_normal.py	10548	0644	edit dl rm
negative_binomial.py	4091	0644	edit dl rm
normal.py	3351	0644	edit dl rm
one_hot_categorical.py	4375	0644	edit dl rm
pareto.py	2057	0644	edit dl rm
poisson.py	2066	0644	edit dl rm
relaxed_bernoulli.py	5360	0644	edit dl rm
relaxed_categorical.py	5202	0644	edit dl rm
studentT.py	3550	0644	edit dl rm
transformed_distribution.py	8270	0644	edit dl rm
transforms.py	38408	0644	edit dl rm
uniform.py	3112	0644	edit dl rm
utils.py	6196	0644	edit dl rm
von_mises.py	5091	0644	edit dl rm
weibull.py	2854	0644	edit dl rm
__init__.py	5884	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/distributions/lkj_cholesky.py (6124B)

"""
This closely follows the implementation in NumPyro (https://github.com/pyro-ppl/numpyro).

Original copyright notice:

# Copyright: Contributors to the Pyro project.
# SPDX-License-Identifier: Apache-2.0
"""

import math

import torch
from torch.distributions import constraints, Beta
from torch.distributions.distribution import Distribution
from torch.distributions.utils import broadcast_all

class LKJCholesky(Distribution):
    r"""
    LKJ distribution for lower Cholesky factor of correlation matrices.
    The distribution is controlled by ``concentration`` parameter :math:`\eta`
    to make the probability of the correlation matrix :math:`M` generated from
    a Cholesky factor propotional to :math:`\det(M)^{\eta - 1}`. Because of that,
    when ``concentration == 1``, we have a uniform distribution over Cholesky
    factors of correlation matrices. Note that this distribution samples the
    Cholesky factor of correlation matrices and not the correlation matrices
    themselves and thereby differs slightly from the derivations in [1] for
    the `LKJCorr` distribution. For sampling, this uses the Onion method from
    [1] Section 3.

L ~ LKJCholesky(dim, concentration)
        X = L @ L' ~ LKJCorr(dim, concentration)

Example::

>>> l = LKJCholesky(3, 0.5)
        >>> l.sample()  # l @ l.T is a sample of a correlation 3x3 matrix
        tensor([[ 1.0000,  0.0000,  0.0000],
                [ 0.3516,  0.9361,  0.0000],
                [-0.1899,  0.4748,  0.8593]])

Args:
        dimension (dim): dimension of the matrices
        concentration (float or Tensor): concentration/shape parameter of the
            distribution (often referred to as eta)

**References**

[1] `Generating random correlation matrices based on vines and extended onion method`,
    Daniel Lewandowski, Dorota Kurowicka, Harry Joe.
    """
    arg_constraints = {'concentration': constraints.positive}
    support = constraints.corr_cholesky

def __init__(self, dim, concentration=1., validate_args=None):
        if dim < 2:
            raise ValueError(f'Expected dim to be an integer greater than or equal to 2. Found dim={dim}.')
        self.dim = dim
        self.concentration, = broadcast_all(concentration)
        batch_shape = self.concentration.size()
        event_shape = torch.Size((dim, dim))
        # This is used to draw vectorized samples from the beta distribution in Sec. 3.2 of [1].
        marginal_conc = self.concentration + 0.5 * (self.dim - 2)
        offset = torch.arange(self.dim - 1, dtype=self.concentration.dtype, device=self.concentration.device)
        offset = torch.cat([offset.new_zeros((1,)), offset])
        beta_conc1 = offset + 0.5
        beta_conc0 = marginal_conc.unsqueeze(-1) - 0.5 * offset
        self._beta = Beta(beta_conc1, beta_conc0)
        super(LKJCholesky, self).__init__(batch_shape, event_shape, validate_args)

def expand(self, batch_shape, _instance=None):
        new = self._get_checked_instance(LKJCholesky, _instance)
        batch_shape = torch.Size(batch_shape)
        new.dim = self.dim
        new.concentration = self.concentration.expand(batch_shape)
        new._beta = self._beta.expand(batch_shape + (self.dim,))
        super(LKJCholesky, new).__init__(batch_shape, self.event_shape, validate_args=False)
        new._validate_args = self._validate_args
        return new

def sample(self, sample_shape=torch.Size()):
        # This uses the Onion method, but there are a few differences from [1] Sec. 3.2:
        # - This vectorizes the for loop and also works for heterogeneous eta.
        # - Same algorithm generalizes to n=1.
        # - The procedure is simplified since we are sampling the cholesky factor of
        #   the correlation matrix instead of the correlation matrix itself. As such,
        #   we only need to generate `w`.
        y = self._beta.sample(sample_shape).unsqueeze(-1)
        u_normal = torch.randn(self._extended_shape(sample_shape),
                               dtype=y.dtype,
                               device=y.device).tril(-1)
        u_hypersphere = u_normal / u_normal.norm(dim=-1, keepdim=True)
        # Replace NaNs in first row
        u_hypersphere[..., 0, :].fill_(0.)
        w = torch.sqrt(y) * u_hypersphere
        # Fill diagonal elements; clamp for numerical stability
        eps = torch.finfo(w.dtype).tiny
        diag_elems = torch.clamp(1 - torch.sum(w**2, dim=-1), min=eps).sqrt()
        w += torch.diag_embed(diag_elems)
        return w

def log_prob(self, value):
        # See: https://mc-stan.org/docs/2_25/functions-reference/cholesky-lkj-correlation-distribution.html
        # The probability of a correlation matrix is proportional to
        #   determinant ** (concentration - 1) = prod(L_ii ^ 2(concentration - 1))
        # Additionally, the Jacobian of the transformation from Cholesky factor to
        # correlation matrix is:
        #   prod(L_ii ^ (D - i))
        # So the probability of a Cholesky factor is propotional to
        #   prod(L_ii ^ (2 * concentration - 2 + D - i)) = prod(L_ii ^ order_i)
        # with order_i = 2 * concentration - 2 + D - i
        if self._validate_args:
            self._validate_sample(value)
        diag_elems = value.diagonal(dim1=-1, dim2=-2)[..., 1:]
        order = torch.arange(2, self.dim + 1)
        order = 2 * (self.concentration - 1).unsqueeze(-1) + self.dim - order
        unnormalized_log_pdf = torch.sum(order * diag_elems.log(), dim=-1)
        # Compute normalization constant (page 1999 of [1])
        dm1 = self.dim - 1
        alpha = self.concentration + 0.5 * dm1
        denominator = torch.lgamma(alpha) * dm1
        numerator = torch.mvlgamma(alpha - 0.5, dm1)
        # pi_constant in [1] is D * (D - 1) / 4 * log(pi)
        # pi_constant in multigammaln is (D - 1) * (D - 2) / 4 * log(pi)
        # hence, we need to add a pi_constant = (D - 1) * log(pi) / 2
        pi_constant = 0.5 * dm1 * math.log(math.pi)
        normalize_term = pi_constant + numerator - denominator
        return unnormalized_log_pdf - normalize_term

cmd: