(coords, radius, L=8)
| 513 | |
| 514 | |
| 515 | def IntegratedPositionalEncoding(coords, radius, L=8): |
| 516 | |
| 517 | # adapted from mipnerf https://github.com/google/mipnerf |
| 518 | def expected_sin(x, x_var): |
| 519 | """Estimates mean and variance of sin(z), z ~ N(x, var).""" |
| 520 | |
| 521 | # When the variance is wide, shrink sin towards zero. |
| 522 | y = torch.exp(-0.5 * x_var) * torch.sin(x) |
| 523 | y_var = torch.clip(0.5 * (1 - torch.exp(-2 * x_var) * torch.cos(2 * x)) - y**2, 0) |
| 524 | return y, y_var |
| 525 | |
| 526 | def integrated_pos_enc(x_coord, min_deg, max_deg): |
| 527 | """Encode `x` with sinusoids scaled by 2^[min_deg:max_deg-1].""" |
| 528 | |
| 529 | x, x_cov_diag = x_coord |
| 530 | scales = torch.tensor([2**i for i in range(int(min_deg), int(max_deg))], device=x.device) |
| 531 | shape = list(x.shape[:-1]) + [-1] |
| 532 | |
| 533 | y = torch.reshape(x[..., None, :] * scales[:, None], shape) |
| 534 | y_var = torch.reshape(x_cov_diag[..., None, :] * scales[:, None]**2, shape) |
| 535 | |
| 536 | return expected_sin( |
| 537 | torch.cat([y, y + 0.5 * np.pi], dim=-1), |
| 538 | torch.cat([y_var] * 2, dim=-1))[0] |
| 539 | |
| 540 | means = coords |
| 541 | covs = (radius**2 / 4) * torch.ones((1, 2), device=coords.device).repeat(coords.shape[-2], 1) |
| 542 | return integrated_pos_enc((means, covs), 0, L) |
| 543 | |
| 544 | |
| 545 | class FFPositionalEncoding(nn.Module): |
nothing calls this directly
no test coverage detected