Sample spherical gaussian distributions. Args: sample_shape: (*sample_shape). Can be an empty list. Returns: samples: (*sample_shape, *m_shape, 3) Sampling strategy: We are going to first sample assuming u = (0,0,1) for all com
(self, sample_shape: T.List[int])
| 92 | return nlls |
| 93 | |
| 94 | def sample(self, sample_shape: T.List[int]): |
| 95 | """ |
| 96 | Sample spherical gaussian distributions. |
| 97 | |
| 98 | Args: |
| 99 | sample_shape: (*sample_shape). Can be an empty list. |
| 100 | |
| 101 | Returns: |
| 102 | samples: (*sample_shape, *m_shape, 3) |
| 103 | |
| 104 | Sampling strategy: |
| 105 | We are going to first sample assuming u = (0,0,1) for all components |
| 106 | but with the actual k. Then we will rotate the samples by a rotation |
| 107 | matrix. |
| 108 | """ |
| 109 | |
| 110 | # sample assuming u = (0,0,1) |
| 111 | thetas = torch.rand(*sample_shape, *self.m_shape) * 2 * torch.pi # (S, M) |
| 112 | vs = torch.stack((torch.cos(thetas), torch.sin(thetas)), dim=-1) # (S, M, 2) |
| 113 | |
| 114 | etas = torch.rand(*sample_shape, *self.m_shape) # (S, M) |
| 115 | ks = self.k.expand(*sample_shape, *self.m_shape) # (S, M) |
| 116 | ws = 1 + ks.pow(-1) * torch.log(etas + (1 - etas) * torch.exp(-2 * ks)) # (S, M) |
| 117 | ws = ws.unsqueeze(-1) # (S, M, 1) |
| 118 | |
| 119 | samples = torch.cat( |
| 120 | ( |
| 121 | (1 - ws.pow(2)).sqrt() * vs, |
| 122 | ws, |
| 123 | ), dim=-1) # (S, M, 3) |
| 124 | |
| 125 | # now that we have samples for u = (0,0,1), we will rotate the samples |
| 126 | # ori_us = torch.zeros(*samples, *self.m_shape, 3) |
| 127 | # ori_us[..., 2] = 1 |
| 128 | # new_us = self.u.expand(*samples, *self.m_shape, 3) # (S, M, 3) |
| 129 | # Rs = rigid_motion.get_min_R( |
| 130 | # ori_us, |
| 131 | # new_us, |
| 132 | # ) # (S, M, 3, 3) new_us = Rs @ ori_us |
| 133 | |
| 134 | # we construct a rotation matrix that rotates (0,0,1) to u. |
| 135 | # it might not be the geodestic rotation matrix, but it is ok |
| 136 | # since spherical gaussian is symmetric around the mean direction. |
| 137 | ys = torch.zeros(*sample_shape, *self.m_shape, 3) |
| 138 | ys[..., 1] = 1 |
| 139 | Rs = rigid_motion.construct_coord_frame( |
| 140 | z=self.u.expand(*sample_shape, *self.m_shape, 3), # (S, M, 3) |
| 141 | y=ys, # (S, M, 3) |
| 142 | ) # (S, M, 3, 3) |
| 143 | |
| 144 | samples = (Rs @ samples.unsqueeze(-1)).squeeze(-1) # (S, M, 3) |
| 145 | |
| 146 | return samples |