Function eval_shfs_4d

utils/sh_utils.py:115–223

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, deg_t, sh, dirs, dirs_t, l = torch.pi)

    return result

def eval_shfs_4d(deg, deg_t, sh, dirs, dirs_t, l = torch.pi):
    """
    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.ndarray unit directions [..., 3]
    Returns:
        [..., C]
    """
    assert deg <= 4 and deg >= 0
    # coeff = (deg + 1) ** 2
    # assert sh.shape[-1] >= coeff

    l0m0 = C0
    result = l0m0 * sh[..., 0]
    
    if deg > 0:
        x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]
        l1m1 = -1 * C1 * y
        l1m0 = C1 * z
        l1p1 = -1 * C1 * x
        
        result = (result + 
                  l1m1 * sh[..., 1] +
                  l1m0 * sh[..., 2] +
                  l1p1 * sh[..., 3])

        if deg > 1:
            xx, yy, zz = x * x, y * y, z * z
            xy, yz, xz = x * y, y * z, x * z
            
            l2m2 = C2[0] * xy
            l2m1 = C2[1] * yz
            l2m0 = C2[2] * (2.0 * zz - xx - yy)
            l2p1 = C2[3] * xz
            l2p2 = C2[4] * (xx - yy)
            
            result = (result + 
                  l2m2 * sh[..., 4] +
                  l2m1 * sh[..., 5] +
                  l2m0 * sh[..., 6] +
                  l2p1 * sh[..., 7] +
                  l2p2 * sh[..., 8])

            if deg > 2:
                l3m3 = C3[0] * y * (3 * xx - yy)
                l3m2 = C3[1] * xy * z
                l3m1 = C3[2] * y * (4 * zz - xx - yy)
                l3m0 = C3[3] * z * (2 * zz - 3 * xx - 3 * yy)
                l3p1 = C3[4] * x * (4 * zz - xx - yy)
                l3p2 = C3[5] * z * (xx - yy)
                l3p3 = C3[6] * x * (xx - 3 * yy)
                
                result = (result +