Table of Contents
Shortcuts

torch.distributions.lkj_cholesky 的源代码

"""
这紧密遵循NumPyro中的实现（https://github.com/pyro-ppl/numpyro）。

原始版权声明：

# 版权：Pyro项目的贡献者。
# SPDX许可证标识符：Apache-2.0
"""

import math

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

__all__ = ["LKJCholesky"]



[docs]class LKJCholesky(Distribution):
    r"""
    LKJ分布用于相关矩阵的下三角Cholesky因子。
    该分布由``concentration``参数 :math:`\eta` 控制，
    以使从Cholesky因子生成的相关矩阵 :math:`M` 的概率与 :math:`\det(M)^{\eta - 1}` 成正比。
    因此，当 ``concentration == 1`` 时，我们在相关矩阵的Cholesky因子上有一个均匀分布::

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

    请注意，此分布采样的是相关矩阵的Cholesky因子，而不是相关矩阵本身，
    因此与[1]中`LKJCorr`分布的推导略有不同。
    对于采样，此方法使用了[1]第3节中的Onion方法。

    示例::

        >>> # xdoctest: +IGNORE_WANT("非确定性")
        >>> l = LKJCholesky(3, 0.5)
        >>> l.sample()  # l @ l.T 是3x3相关矩阵的一个样本
        tensor([[ 1.0000,  0.0000,  0.0000],
                [ 0.3516,  0.9361,  0.0000],
                [-0.1899,  0.4748,  0.8593]])

    参数:
        dimension (dim): 矩阵的维度
        concentration (float or Tensor): 分布的浓度/形状参数（通常称为 eta）

    **参考文献**

    [1] `基于藤蔓和扩展洋葱方法生成随机相关矩阵` (2009),
    Daniel Lewandowski, Dorota Kurowicka, Harry Joe.
    多元分析杂志。100。10.1016/j.jmva.2009.04.008
    """
    arg_constraints = {"concentration": constraints.positive}
    support = constraints.corr_cholesky

    def __init__(self, dim, concentration=1.0, 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))
        # 这用于从[1]第3.2节中的beta分布中绘制向量化样本。
        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().__init__(batch_shape, event_shape, validate_args)


[docs]    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



[docs]    def sample(self, sample_shape=torch.Size()):
        # 这使用了Onion方法，但与[1]第3.2节有一些不同：
        # - 这向量化了for循环，并且适用于异构eta。
        # - 相同的算法推广到n=1。
        # - 该过程被简化，因为我们采样的是相关矩阵的Cholesky因子，而不是相关矩阵本身。
        #   因此，我们只需要生成`w`。
        y = self.<span
优云智算