https://github.com/cheind/torch-spherical-harmonics
Real Spherical Harmonics for PyTorch
https://github.com/cheind/torch-spherical-harmonics
pytorch spherical-harmonics symbolic
Last synced: 2 months ago
JSON representation
Real Spherical Harmonics for PyTorch
- Host: GitHub
- URL: https://github.com/cheind/torch-spherical-harmonics
- Owner: cheind
- License: mit
- Created: 2022-11-03T18:15:41.000Z (almost 3 years ago)
- Default Branch: main
- Last Pushed: 2022-12-05T12:39:50.000Z (almost 3 years ago)
- Last Synced: 2025-04-14T00:26:49.215Z (6 months ago)
- Topics: pytorch, spherical-harmonics, symbolic
- Language: Python
- Homepage:
- Size: 2.82 MB
- Stars: 22
- Watchers: 2
- Forks: 2
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
## **torch-spherical-harmonics**
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]$.
## Properties
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._
## Usage
To compute all $Y_{nm}$ values up to degree 3 use:
```python
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`](./torchsh/rsh.py) in your project, which requires only `torch` to be installed.
## Code Generation
We use `sympy` to generate RSH expressions in Cartesian form reyling on the [Herglotzian](https://en.wikipedia.org/wiki/Spherical_harmonics#Separated_Cartesian_form) 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
```
## References
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.