Class SphericalGaussian

plib/spherical_gaussian.py:12–146

r""" This class implements the spherical Gaussian distribution (i.e., the von Mises-Fisher distribution in 3D). Let :math:`w` in :math:`\mathbb{S}^2` and :math:`k \ge 0`, the spherical Gaussian distribution is defined as .. math:: f(w) & = \frac{1}{4 \pi}, if k = 0 \\

class SphericalGaussian:
    r"""
    This class implements the spherical Gaussian distribution (i.e., the von Mises-Fisher distribution in 3D).

    Let :math:`w` in :math:`\mathbb{S}^2` and :math:`k \ge 0`, the spherical Gaussian distribution is defined as

    .. math::
        f(w) & = \frac{1}{4 \pi},  if k = 0   \\
             & = \frac{k}{2 \pi (1 - exp(-2 k))} exp(k (u^T w - 1)),  k > 0

    Ref: http://www.mitsuba-renderer.org/~wenzel/files/vmf.pdf
    """

    def __init__(
            self,
            u: torch.Tensor,
            k: T.Union[torch.Tensor, None],
            log_k: torch.Tensor = None,
            normalize_u: bool = False,
    ):
        """
        Args:
            u: (*, 3)
                the mean direction.
                (*) is the dimension of independent components`.
                We are going to use M to represent it for simplicity.
            k: (*,) or (,)
                the concentration on the sphere  (large k -> more concentrated)
            log_k: (*,) or (,)
                log of k, can be None. If given, k will be ignored.
        """
        *m_shape, dim = u.shape
        self.m_shape = m_shape
        self.normalize_u = normalize_u

        assert u.size(-1) == 3
        if self.normalize_u:
            u = torch.nn.functional.normalize(u, dim=-1)

        self.u = u

        if isinstance(k, (float, int)):
            k = torch.ones(*m_shape, device=u.device) * k

        # assert k.shape == m_shape, f'{k.shape} {m_shape}'
        self.k = k
        if log_k is None:
            self.log_k = self.k.log()
        else:
            self.log_k = log_k
            self.k = self.log_k.exp()

        self.log_2pi = math.log(2 * math.pi)

    def compute_neg_log_likelihoods(
            self,
            samples: torch.Tensor,
            normalize: bool = True,