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
(
self,
u: torch.Tensor,
k: T.Union[torch.Tensor, None],
log_k: torch.Tensor = None,
normalize_u: bool = False,
)
| 23 | """ |
| 24 | |
| 25 | def __init__( |
| 26 | self, |
| 27 | u: torch.Tensor, |
| 28 | k: T.Union[torch.Tensor, None], |
| 29 | log_k: torch.Tensor = None, |
| 30 | normalize_u: bool = False, |
| 31 | ): |
| 32 | """ |
| 33 | Args: |
| 34 | u: (*, 3) |
| 35 | the mean direction. |
| 36 | (*) is the dimension of independent components`. |
| 37 | We are going to use M to represent it for simplicity. |
| 38 | k: (*,) or (,) |
| 39 | the concentration on the sphere (large k -> more concentrated) |
| 40 | log_k: (*,) or (,) |
| 41 | log of k, can be None. If given, k will be ignored. |
| 42 | """ |
| 43 | *m_shape, dim = u.shape |
| 44 | self.m_shape = m_shape |
| 45 | self.normalize_u = normalize_u |
| 46 | |
| 47 | assert u.size(-1) == 3 |
| 48 | if self.normalize_u: |
| 49 | u = torch.nn.functional.normalize(u, dim=-1) |
| 50 | |
| 51 | self.u = u |
| 52 | |
| 53 | if isinstance(k, (float, int)): |
| 54 | k = torch.ones(*m_shape, device=u.device) * k |
| 55 | |
| 56 | # assert k.shape == m_shape, f'{k.shape} {m_shape}' |
| 57 | self.k = k |
| 58 | if log_k is None: |
| 59 | self.log_k = self.k.log() |
| 60 | else: |
| 61 | self.log_k = log_k |
| 62 | self.k = self.log_k.exp() |
| 63 | |
| 64 | self.log_2pi = math.log(2 * math.pi) |
| 65 | |
| 66 | def compute_neg_log_likelihoods( |
| 67 | self, |