| 203 | |
| 204 | # Hierarchical sampling (section 5.2) |
| 205 | def sample_pdf(bins, weights, N_samples, det=False, pytest=False): |
| 206 | # Get pdf |
| 207 | weights = weights + 1e-5 # prevent nans |
| 208 | pdf = weights / torch.sum(weights, -1, keepdim=True) |
| 209 | cdf = torch.cumsum(pdf, -1) |
| 210 | cdf = torch.cat([torch.zeros_like(cdf[...,:1]), cdf], -1) # (batch, len(bins)) |
| 211 | |
| 212 | # Take uniform samples |
| 213 | if det: |
| 214 | u = torch.linspace(0., 1., steps=N_samples) |
| 215 | u = u.expand(list(cdf.shape[:-1]) + [N_samples]) |
| 216 | else: |
| 217 | u = torch.rand(list(cdf.shape[:-1]) + [N_samples]) |
| 218 | |
| 219 | # Pytest, overwrite u with numpy's fixed random numbers |
| 220 | if pytest: |
| 221 | np.random.seed(0) |
| 222 | new_shape = list(cdf.shape[:-1]) + [N_samples] |
| 223 | if det: |
| 224 | u = np.linspace(0., 1., N_samples) |
| 225 | u = np.broadcast_to(u, new_shape) |
| 226 | else: |
| 227 | u = np.random.rand(*new_shape) |
| 228 | u = torch.Tensor(u) |
| 229 | |
| 230 | # Invert CDF |
| 231 | u = u.contiguous() |
| 232 | # inds = searchsorted(cdf, u, side='right') |
| 233 | inds = torch.searchsorted(cdf, u, right=True) |
| 234 | below = torch.max(torch.zeros_like(inds-1), inds-1) |
| 235 | above = torch.min((cdf.shape[-1]-1) * torch.ones_like(inds), inds) |
| 236 | inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2) |
| 237 | |
| 238 | # cdf_g = tf.gather(cdf, inds_g, axis=-1, batch_dims=len(inds_g.shape)-2) |
| 239 | # bins_g = tf.gather(bins, inds_g, axis=-1, batch_dims=len(inds_g.shape)-2) |
| 240 | matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]] |
| 241 | cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g) |
| 242 | bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g) |
| 243 | |
| 244 | denom = (cdf_g[...,1]-cdf_g[...,0]) |
| 245 | denom = torch.where(denom<1e-5, torch.ones_like(denom), denom) |
| 246 | t = (u-cdf_g[...,0])/denom |
| 247 | samples = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0]) |
| 248 | |
| 249 | return samples |