cheind / torch-spherical-harmonics Goto Github PK

Real spherical harmonics (RSH) in Cartesian form for PyTorch. The resulting source code is auto-generated by converting optimized symbolic RSH expressions to PyTorch.

The following plot shows the first real spherical harmonics $Y_{nm}$ of degree $n < 6$ and order $-n \le m \le n$ as a function of polar coordinates $\theta \in [0,\pi]$ and $\phi \in [0,2\pi]$.

The RSH functions provided in this package are ortho-normalized over the surface of the unit sphere. That is

$$ \int\limits_0^\pi\int\limits_0^{2\pi}Y_{nm}(\mathbf r(\theta,\phi))Y_{n'm'}(\mathbf r(\theta,\phi)) \left\lVert \frac{\partial\mathbf{r}}{\partial\theta}\times\frac{\partial\mathbf{r}}{\partial\phi} \right\rVert \thinspace d \phi \thinspace d \theta = \delta_{nn'}\delta_{mm'}, $$

with

$$ \mathbf r(\theta,\phi) = \sin\theta\cos\phi\mathbf{\hat i} + \sin\theta\sin\phi\mathbf{\hat j} + \cos\theta\mathbf{\hat k}, $$

and

$$ \left\lVert \frac{\partial\mathbf{r}}{\partial\theta} \times \frac{\partial\mathbf{r}}{\partial\phi} \right\lVert = \sin\theta. $$

Note that due to some rendering issue, you might not see the double vertical bars around the cross product.

To compute all $Y_{nm}$ values up to degree 3 use:

import torch
from torchsh import rsh_cart_3

xyz = ...             # tensor (N,...,3) of unit-sphere points
ynm = rsh_cart_3(xyz) # tensor (N,...,16) of Ynm values
                      # with Ynm at index `n*(n+1) + m`

This relies on pre-generated RSH functions, which torchsh contains up to degree 8. They all follow the same naming convention rsh_cart_{degree}.

If you do not want to add a new library dependency, you may just as well just include torchsh/rsh.py in your project, which requires only torch to be installed.

We use sympy to generate RSH expressions in Cartesian form reyling on the Herglotzian definition. These expressions are simplified and transformed into Python/PyTorch functions using a code template and the string engine mako. In the tests we use sympy.Znm to verify our numerical results.

We initially intended to use sympy.Znm directly for code generation, by substituting polar coordinate definitions with respective Cartesian ones, but found that substitution did not work all of the times. This is the reason we switched to a different generating definition.

Code can be generated via

$ python -m torchsh.symbolic.codegen --help

which requires all dev-requirements.txt to be installed. To run the unit tests call

$ pytest

The basic idea for using sympy to generate code for RSH functions is taken from https://nvlabs.github.io/instant-ngp, where it is used to generate Cuda code for forward and backward passes.