most changes for v0.6 release

AUTHORS

 The code was authored by the following people:
 
 Boris Bonev - NVIDIA Corporation
-Christian Hundt - NVIDIA Corporation
 Thorsten Kurth - NVIDIA Corporation
+Christian Hundt - NVIDIA Corporation
 Nikola Kovachki - NVIDIA Corporation
+Jean Kossaifi - NVIDIA Corporation
\ No newline at end of file

Changelog.md

@@ -2,18 +2,23 @@
 
 ## Versioning
 
-### v0.5
+### v0.6.0
+
+* Added SFNO example
+* Cleanup of the repository and added PyPI
+
+### v0.5.0
 
 * Reworked distributed SHT
 * Module for sampling Gaussian Random Fields on the sphere
 
-### v0.4
+### v0.4.0
 
 * Computation of associated Legendre polynomials
     * changed algorithm to compute the associated Legendre polynomials for improved stability
 * Improved Readme
 
-### v0.3
+### v0.3.0
 
 * Vector Spherical Harmonic Transforms
     * projects vector-valued fields onto the vector Spherical Harmonics
@@ -28,13 +33,13 @@
 * Unit tests
 * Refactor of the API
 
-### v0.2
+### v0.2.0
 
 * Renaming from torch_sht to torch_harmonics
 * Adding distributed SHT support
 * New logo
 
-### v0.1
+### v0.1.0
 
 * Single GPU forward and backward transform
-* Minimal code example and notebook 
+* Minimal code example and notebook
\ No newline at end of file

README.md

@@ -36,13 +36,11 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 <!-- ## What is torch-harmonics? -->
 
-Spherical Harmonic Transforms (SHTs) are the counterpart to Fourier transforms on the sphere. As such they are an invaluable tool for signal-processing on the sphere.
-
-`torch_harmonics` is a differentiable implementation of the SHT in PyTorch. 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.
+`torch_harmonics` is a differentiable implementation of the Spherical Harmonic transform in PyTorch. 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.
 
 `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.
 
-`torch_harmonics` has been used to implement a variety of differentiable PDE solvers which generated the animations below.
+`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).
 
 
 <table border="0" cellspacing="0" cellpadding="0">
@@ -83,9 +81,10 @@ docker run --gpus all -it --rm --ipc=host --ulimit memlock=-1 --ulimit stack=671
  - Christian Hundt (chundt@nvidia.com)
  - Thorsten Kurth (tkurth@nvidia.com)
  - Nikola Kovachki (nkovachki@nvidia.com)
+ - Jean Kossaifi (jkossaifi@nvidia.com)
 
 ## Implementation
-The implementation follows the paper "Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations", N. Schaeffer, G3: Geochemistry, Geophysics, Geosystems. 
+The implementation follows the algorithm as presented in [2].
 
 ### Spherical harmonic transform
 
@@ -150,14 +149,34 @@ coeffs = sht(signal)
 
 `torch_harmonics` also implements a distributed variant of the SHT located in `torch-harmonics.distributed`.
 
+### Cite us
+
+If you use `torch-harmonics` in an academic paper, please cite [1]
+
+```
+@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}
+}
+```
+
 ## References
 
 <a id="1">[1]</a> 
-Schaeffer, N.;
+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.;
 Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations;
 G3: Geochemistry, Geophysics, Geosystems, 2013.
 
-<a id="1">[2]</a> 
-Wang, B., Wang, L., Xie, Z.;
+<a id="1">[3]</a> 
+Wang B., Wang L., Xie Z.;
 Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids;
 Adv Comput Math, 2018.

examples/shallow_water_equations.py

@@ -365,7 +365,7 @@ class ShallowWaterSolver(nn.Module):
             if ccrs is None:
                 raise ImportError("Couldn't import Cartopy")
 
-            proj = ccrs.Orthographic(central_longitude=0.0, central_latitude=45.0)
+            proj = ccrs.Orthographic(central_longitude=0.0, central_latitude=-45.0)
 
             #ax = plt.gca(projection=proj, frameon=True)
             ax = fig.add_subplot(projection=proj)

notebooks/quadrature.ipynb

@@ -531,7 +531,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.6"
+   "version": "3.10.6 (main, Nov 14 2022, 16:10:14) [GCC 11.3.0]"
   },
   "orig_nbformat": 4,
   "vscode": {

setup.py

@@ -28,19 +28,74 @@
 # 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.
 #
+try:
+    from setuptools import setup, find_packages
+except ImportError:
+    from distutils.core import setup, find_packages
 
-from setuptools import setup
-
-setup(
-   name='torch_harmonics',
-   version='0.5',
-   author='Boris Bonev',
-   author_email='bbonev@nvidia.com',
-   packages=['torch_harmonics','torch_harmonics.distributed'],
-   scripts=[],
-   url='https://github.com/NVIDIA/torch-harmonics',
-   license='LICENSE.md',
-   description='a differentiable spherical harmonic transform for PyTorch',
-   long_description=open('README.md').read(),
-   install_requires=['torch','numpy']
-)
+import re
+from pathlib import Path
+
+
+def version(root_path):
+    """Returns the version taken from __init__.py
+
+    Parameters
+    ----------
+    root_path : pathlib.Path
+        path to the root of the package
+
+    Reference
+    ---------
+    https://packaging.python.org/guides/single-sourcing-package-version/
+    """
+    version_path = root_path.joinpath('neuralop', '__init__.py')
+    with version_path.open() as f:
+        version_file = f.read()
+    version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
+                              version_file, re.M)
+    if version_match:
+        return version_match.group(1)
+    raise RuntimeError("Unable to find version string.")
+
+
+def readme(root_path):
+    """Returns the text content of the README.md of the package
+
+    Parameters
+    ----------
+    root_path : pathlib.Path
+        path to the root of the package
+    """
+    with root_path.joinpath('README.md').open(encoding='UTF-8') as f:
+        return f.read()
+
+
+root_path = Path(__file__).parent
+README = readme(root_path)
+VERSION = version(root_path)
+
+
+config = {
+    'name': 'torch_harmonics',
+    'packages': find_packages(),
+    'description': 'A differentiable spherical harmonic transform for PyTorch.',
+    'long_description': README,
+    'long_description_content_type' : 'text/x-rst',
+    'url' : 'https://github.com/NVIDIA/torch-harmonics',
+    'authors': [
+        {'name': "Boris Bonev", 'email': "bbonev@nvidia.com"},
+        ],
+    'version': VERSION,
+    'install_requires': ['torch','numpy'],
+    'license': 'Modified BSD',
+    'scripts': [],
+    'include_package_data': True,
+    'classifiers': [
+        'Topic :: Scientific/Engineering',
+        'License :: OSI Approved :: BSD License',
+        'Programming Language :: Python :: 3'
+    ],
+}
+
+setup(**config)

torch_harmonics/__init__.py

@@ -29,6 +29,8 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #
 
+__version__ = '0.6.0'
+
 from .sht import RealSHT, InverseRealSHT, RealVectorSHT, InverseRealVectorSHT
 from . import quadrature
 from . import random_fields