README.md 9.75 KB
Newer Older
Boris Bonev's avatar
Boris Bonev committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
<!-- 
SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.

SPDX-License-Identifier: BSD-3-Clause
 
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this
   list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its
   contributors may be used to endorse or promote products derived from
   this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-->

Boris Bonev's avatar
Boris Bonev committed
31
<p align="center">
Boris Bonev's avatar
Boris Bonev committed
32
    <img src="https://raw.githubusercontent.com/NVIDIA/torch-harmonics/main/images/logo/logo.png"  width="568">
Boris Bonev's avatar
Boris Bonev committed
33
    <br>
Boris Bonev's avatar
Boris Bonev committed
34
    <a href="https://github.com/NVIDIA/torch-harmonics/actions/workflows/tests.yml"><img src="https://github.com/NVIDIA/torch-harmonics/actions/workflows/tests.yml/badge.svg"></a>
Boris Bonev's avatar
Boris Bonev committed
35
    <a href="https://pypi.org/project/torch_harmonics/"><img src="https://img.shields.io/pypi/v/torch_harmonics"></a>
Boris Bonev's avatar
Boris Bonev committed
36
37
</p>

38
39
[**<p style="text-align: center;">Overview**](#overview) | [**Installation**](#installation) | [**More information**](#more-about-torch-harmonics) | [**Getting started**](#getting-started) | [**Contributors**](#contributors) | [**Cite us**](#cite-us) | [**References</p>**](#references)

Boris Bonev's avatar
Boris Bonev committed
40
41
42
43
<!--
[![pypi](https://img.shields.io/pypi/v/torch_harmonics)](https://pypi.org/project/torch_harmonics/)
-->

Boris Bonev's avatar
Boris Bonev committed
44
45
<!-- # torch-harmonics: differentiable harmonic transforms -->

46
## Overview
Boris Bonev's avatar
Boris Bonev committed
47

Boris Bonev's avatar
Boris Bonev committed
48
torch-harmonics is a differentiable implementation of the Spherical Harmonic transform in PyTorch. It was originally implemented to enable Spherical Fourier Neural Operators (SFNO). It uses quadrature rules to compute the projection onto the associated Legendre polynomials and FFTs for the projection onto the harmonic basis. This algorithm tends to outperform others with better asymptotic scaling for most practical purposes.
Boris Bonev's avatar
Boris Bonev committed
49

Boris Bonev's avatar
Boris Bonev committed
50
torch-harmonics uses PyTorch primitives to implement these operations, making it fully differentiable. Moreover, the quadrature can be distributed onto multiple ranks making it spatially distributed.
Boris Bonev's avatar
Boris Bonev committed
51

Boris Bonev's avatar
Boris Bonev committed
52
torch-harmonics has been used to implement a variety of differentiable PDE solvers which generated the animations below. Moreover, it has enabled the development of Spherical Fourier Neural Operators (SFNOs) [1].
Boris Bonev's avatar
Boris Bonev committed
53
54
55
56


<table border="0" cellspacing="0" cellpadding="0">
    <tr>
Boris Bonev's avatar
Boris Bonev committed
57
58
59
        <td><img src="https://media.githubusercontent.com/media/NVIDIA/torch-harmonics/main/images/sfno.gif"  width="240"></td>
        <td><img src="https://media.githubusercontent.com/media/NVIDIA/torch-harmonics/main/images/zonal_jet.gif"  width="240"></td>
        <td><img src="https://media.githubusercontent.com/media/NVIDIA/torch-harmonics/main/images/allen-cahn.gif"  width="240"></td>
60
    </tr>
Boris Bonev's avatar
Boris Bonev committed
61
<!--     <tr>
Boris Bonev's avatar
Boris Bonev committed
62
63
64
        <td style="text-align:center; border-style : hidden!important;">Shallow Water Eqns.</td>
        <td style="text-align:center; border-style : hidden!important;">Ginzburg-Landau Eqn.</td>
        <td style="text-align:center; border-style : hidden!important;">Allen-Cahn Eqn.</td>
Boris Bonev's avatar
Boris Bonev committed
65
    </tr>  -->
Boris Bonev's avatar
Boris Bonev committed
66
67
68
69
</table>


## Installation
Boris Bonev's avatar
Boris Bonev committed
70
71
Download directyly from PyPI:

72
```bash
Boris Bonev's avatar
Boris Bonev committed
73
74
75
pip install torch-harmonics
```

Boris Bonev's avatar
Boris Bonev committed
76
77
Build in your environment using the Python package:

78
```bash
Boris Bonev's avatar
Boris Bonev committed
79
git clone git@github.com:NVIDIA/torch-harmonics.git
80
81
cd torch-harmonics
pip install -e .
Boris Bonev's avatar
Boris Bonev committed
82
83
84
85
```

Alternatively, use the Dockerfile to build your custom container after cloning:

86
```bash
Boris Bonev's avatar
Boris Bonev committed
87
git clone git@github.com:NVIDIA/torch-harmonics.git
88
cd torch-harmonics
Boris Bonev's avatar
Boris Bonev committed
89
90
91
92
docker build . -t torch_harmonics
docker run --gpus all -it --rm --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 torch_harmonics
```

93
## More about torch-harmonics
Boris Bonev's avatar
Boris Bonev committed
94

95
### Spherical harmonics
Boris Bonev's avatar
Boris Bonev committed
96

97
98
99
100
101
102
103
104
105
The [spherical harmonics](https://en.wikipedia.org/wiki/Spherical_harmonics) are special functions defined on the two-dimensional sphere $S^2$ (embedded in three dimensions). They form an orthonormal basis of the space of square-integrable functions defined on the sphere ($L^2(S^2)$) and are comparable to the harmonic functions defined on a circle/torus. The spherical harmonics are defined as

$$
Y_l^m(\theta, \lambda) = \sqrt{\frac{(2l + 1)}{4 \pi} \frac{(l - m)!}{(l + m)!}} \, P_l^m(\cos \theta) \, \exp(im\lambda),
$$

where $\theta$ and $\lambda$ are colatitude and longitude respectively, and $P_l^m$ the normalized, [associated Legendre polynomials](https://en.wikipedia.org/wiki/Associated_Legendre_polynomials).

<center><img src="https://media.githubusercontent.com/media/NVIDIA/torch-harmonics/main/images/spherical_harmonics.gif" width="432"></center>
Boris Bonev's avatar
Boris Bonev committed
106
107
108

### Spherical harmonic transform

109
110
111
112
113
114
115
116
117
118
The spherical harmonic transform (SHT)

$$
\tilde{f}_l^m = \int_{S^2}  \overline{Y_l^m}(\theta, \lambda) \, f(\theta, \lambda) \; \mathrm{d} \mu(\theta, \lambda)
$$

realizes the projection of a signal $f(\theta, \lambda)$ on $S^2$ onto the spherical harmonics basis. 

generalizes the Fourier transform on the sphere.

Boris Bonev's avatar
Boris Bonev committed
119
120
121
The truncated series expansion of a function $f$ defined on the surface of a sphere can be written as

$$
122
f(\theta, \lambda) = \sum_{m=-M}^{M} \exp(im\lambda) \sum_{l=|m|}^{M} \tilde{f}_l^m \overline{P_l^m} (\cos \theta),
Boris Bonev's avatar
Boris Bonev committed
123
124
$$

125
where $\theta$ is the colatitude, $\lambda$ the longitude, $\overline{P_l^m}$ the normalized, associated Legendre polynomials and $F_n^m$, the expansion coefficient associated to the mode $(m,n)$.
Boris Bonev's avatar
Boris Bonev committed
126

127
The implementation of the SHT follows the algorithm as presented in [2]. A direct spherical harmonic transform can be accomplished by a Fourier transform
Boris Bonev's avatar
Boris Bonev committed
128
129

$$
130
\tilde{f}^m(\theta) = \frac{1}{2 \pi} \int_{0}^{2\pi} f(\theta, \lambda) \, \exp(-im\lambda)  \; \mathrm{d}\lambda
Boris Bonev's avatar
Boris Bonev committed
131
132
133
134
135
$$

in longitude and a Legendre transform

$$
136
\tilde{f}_l^m = \frac{1}{2} \int_{0}^\pi \tilde{f}^m(\theta) \, \overline{P_l^m} (\cos \theta) \, \sin \theta \;\mathrm{d} \theta
Boris Bonev's avatar
Boris Bonev committed
137
138
139
140
141
142
$$

in latitude.

### Discrete Legendre transform

143
The second integral, which computed the projection onto the Legendre polynomials is realized with quadrature. On the Gaussian grid, we use Gaussian quadrature in the $\cos \theta$ domain. The integral
Boris Bonev's avatar
Boris Bonev committed
144
145

$$
146
\tilde{f}_l^m = \frac{1}{2} \int_{-1}^1 \tilde{f}^m(\arccos x) \, \overline{P_l^m} (x) \; \mathrm{d} x
Boris Bonev's avatar
Boris Bonev committed
147
148
$$

149
is obtained with the substitution $x = \cos \theta$ and then approximated by the sum
Boris Bonev's avatar
Boris Bonev committed
150
151

$$
152
\tilde{f}_l^m = \sum_{j=1}^{N_\theta} \tilde{f}^m(\arccos x_j) \bar{P}_n^m(x_j) \, w_j.
Boris Bonev's avatar
Boris Bonev committed
153
154
$$

155
Here, $x_j \in [-1,1]$ are the quadrature nodes with the respective quadrature weights $w_j$.
Boris Bonev's avatar
Boris Bonev committed
156

157
## Getting started
Boris Bonev's avatar
Boris Bonev committed
158
159
160
161
162

The main functionality of `torch_harmonics` is provided in the form of `torch.nn.Modules` for composability. A minimum example is given by:

```python
import torch
Boris Bonev's avatar
Boris Bonev committed
163
import torch_harmonics as th
Boris Bonev's avatar
Boris Bonev committed
164
165
166
167
168
169
170
171
172

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

nlat = 512
nlon = 2*nlat
batch_size = 32
signal = torch.randn(batch_size, nlat, nlon)

# transform data on an equiangular grid
Boris Bonev's avatar
Boris Bonev committed
173
sht = th.RealSHT(nlat, nlon, grid="equiangular").to(device)
Boris Bonev's avatar
Boris Bonev committed
174
175
176

coeffs = sht(signal)
```
Boris Bonev's avatar
Boris Bonev committed
177

Boris Bonev's avatar
Boris Bonev committed
178
To enable scalable model-parallelism, `torch-harmonics` implements a distributed variant of the SHT located in `torch_harmonics.distributed`.
Boris Bonev's avatar
Boris Bonev committed
179

180
181
182
183
184
185
186
187
188
189
190
191
192
193
Detailed usage of torch-harmonics is demonstrated in a series of notebooks:

1. [Getting started](./notebooks/getting_started.ipynb)
2. [Quadrature](./notebooks/quadrature.ipynb)
3. [Visualizing the spherical harmonics](./notebooks/plot_spherical_harmonics.ipynb)
4. [Solving the Helmholtz equation](./notebooks/helmholtz.ipynb)
5. [Solving the shallow water equations](./notebooks/shallow_water_equations.ipynb)
6. [Training Spherical Fourier Neural Operators](./notebooks/train_sfno.ipynb)

## Contributors

[Boris Bonev](https://bonevbs.github.io) (bbonev@nvidia.com), [Thorsten Kurth](https://github.com/azrael417) (tkurth@nvidia.com), [Christian Hundt](https://github.com/gravitino) (chundt@nvidia.com), [Nikola Kovachki](https://kovachki.github.io) (nkovachki@nvidia.com), [Jean Kossaifi](http://jeankossaifi.com) (jkossaifi@nvidia.com)  

## Cite us
Boris Bonev's avatar
Boris Bonev committed
194
195
196

If you use `torch-harmonics` in an academic paper, please cite [1]

197
```bibtex
Boris Bonev's avatar
Boris Bonev committed
198
199
200
201
202
203
204
205
206
207
@misc{bonev2023spherical,
      title={Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere}, 
      author={Boris Bonev and Thorsten Kurth and Christian Hundt and Jaideep Pathak and Maximilian Baust and Karthik Kashinath and Anima Anandkumar},
      year={2023},
      eprint={2306.03838},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}
```

Boris Bonev's avatar
Boris Bonev committed
208
209
210
## References

<a id="1">[1]</a> 
Boris Bonev's avatar
Boris Bonev committed
211
212
213
214
215
216
Bonev B., Kurth T., Hundt C., Pathak, J., Baust M., Kashinath K., Anandkumar A.;
Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere;
arXiv 2306.0383, 2023.

<a id="1">[2]</a> 
Schaeffer N.;
Boris Bonev's avatar
Boris Bonev committed
217
Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations;
Boris Bonev's avatar
Boris Bonev committed
218
G3: Geochemistry, Geophysics, Geosystems, 2013.
Boris Bonev's avatar
Boris Bonev committed
219

Boris Bonev's avatar
Boris Bonev committed
220
221
<a id="1">[3]</a> 
Wang B., Wang L., Xie Z.;
Boris Bonev's avatar
Boris Bonev committed
222
Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids;
Boris Bonev's avatar
Boris Bonev committed
223
Adv Comput Math, 2018.