Evaluate spherical harmonics at unit directions using hardcoded SH polynomials. Works with torch/np/jnp. ... Can be 0 or more batch dimensions. Args: deg: int SH deg. Currently, 0-3 supported sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] dirs: jnp.nd
(deg, sh, dirs)
| 56 | sh_channels_4d = [1, 6, 16, 33] |
| 57 | |
| 58 | def eval_sh(deg, sh, dirs): |
| 59 | """ |
| 60 | Evaluate spherical harmonics at unit directions |
| 61 | using hardcoded SH polynomials. |
| 62 | Works with torch/np/jnp. |
| 63 | ... Can be 0 or more batch dimensions. |
| 64 | Args: |
| 65 | deg: int SH deg. Currently, 0-3 supported |
| 66 | sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] |
| 67 | dirs: jnp.ndarray unit directions [..., 3] |
| 68 | Returns: |
| 69 | [..., C] |
| 70 | """ |
| 71 | assert deg <= 4 and deg >= 0 |
| 72 | coeff = (deg + 1) ** 2 |
| 73 | assert sh.shape[-1] >= coeff |
| 74 | |
| 75 | result = C0 * sh[..., 0] |
| 76 | if deg > 0: |
| 77 | x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3] |
| 78 | result = (result - |
| 79 | C1 * y * sh[..., 1] + |
| 80 | C1 * z * sh[..., 2] - |
| 81 | C1 * x * sh[..., 3]) |
| 82 | |
| 83 | if deg > 1: |
| 84 | xx, yy, zz = x * x, y * y, z * z |
| 85 | xy, yz, xz = x * y, y * z, x * z |
| 86 | result = (result + |
| 87 | C2[0] * xy * sh[..., 4] + |
| 88 | C2[1] * yz * sh[..., 5] + |
| 89 | C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] + |
| 90 | C2[3] * xz * sh[..., 7] + |
| 91 | C2[4] * (xx - yy) * sh[..., 8]) |
| 92 | |
| 93 | if deg > 2: |
| 94 | result = (result + |
| 95 | C3[0] * y * (3 * xx - yy) * sh[..., 9] + |
| 96 | C3[1] * xy * z * sh[..., 10] + |
| 97 | C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] + |
| 98 | C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] + |
| 99 | C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] + |
| 100 | C3[5] * z * (xx - yy) * sh[..., 14] + |
| 101 | C3[6] * x * (xx - 3 * yy) * sh[..., 15]) |
| 102 | |
| 103 | if deg > 3: |
| 104 | result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] + |
| 105 | C4[1] * yz * (3 * xx - yy) * sh[..., 17] + |
| 106 | C4[2] * xy * (7 * zz - 1) * sh[..., 18] + |
| 107 | C4[3] * yz * (7 * zz - 3) * sh[..., 19] + |
| 108 | C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] + |
| 109 | C4[5] * xz * (7 * zz - 3) * sh[..., 21] + |
| 110 | C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] + |
| 111 | C4[7] * xz * (xx - 3 * yy) * sh[..., 23] + |
| 112 | C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24]) |
| 113 | return result |
| 114 | |
| 115 | def eval_shfs_4d(deg, deg_t, sh, dirs, dirs_t, l = torch.pi): |