Moved SFNO examples into the library to make them available

examples/sfno_examples/train_sfno.py → examples/train_sfno.py

@@ -45,15 +45,13 @@ import pandas as pd
 
 import matplotlib.pyplot as plt
 
+from torch_harmonics.examples.sfno import PdeDataset
+from torch_harmonics.examples.sfno import SphericalFourierNeuralOperatorNet as SFNO
+
 # wandb logging
 import wandb
 wandb.login()
 
-import sys
-sys.path.append(os.path.join(os.path.dirname( __file__), "../"))
-
-from pde_sphere import SphereSolver
-
 def l2loss_sphere(solver, prd, tar, relative=False, squared=False):
     loss = solver.integrate_grid((prd - tar)**2, dimensionless=True).sum(dim=-1)
     if relative:
@@ -166,9 +164,6 @@ def main(train=True, load_checkpoint=False, enable_amp=False):
     if torch.cuda.is_available():
         torch.cuda.set_device(device.index)
 
-    # dataset
-    from utils.pde_dataset import PdeDataset
-
     # 1 hour prediction steps
     dt = 1*3600
     dt_solver = 150
@@ -337,8 +332,6 @@ def main(train=True, load_checkpoint=False, enable_amp=False):
     # from models.unet import UNet
     # models['unet_baseline'] = partial(UNet)
 
-    # SFNO and FNO models
-    from models.sfno import SphericalFourierNeuralOperatorNet as SFNO
     # SFNO models
     models['sfno_sc3_layer4_edim256_linear']    = partial(SFNO, spectral_transform='sht', filter_type='linear', img_size=(nlat, nlon),
                                                      num_layers=4, scale_factor=3, embed_dim=256, operator_type='vector')

examples/sfno_examples/train_sfno.ipynb → notebooks/train_sfno.ipynb

@@ -38,12 +38,6 @@
     "\n",
     "import time\n",
     "\n",
-    "import sys\n",
-    "sys.path.append(\"../../examples\")\n",
-    "\n",
-    "# from torch_harmonics.sht import *\n",
-    "from pde_sphere import SphereSolver\n",
-    "\n",
     "cmap='twilight_shifted'"
    ]
   },
@@ -153,7 +147,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from models.sfno import SphericalFourierNeuralOperatorNet as SFNO"
+    "from torch_harmonics.examples.sfno import SphericalFourierNeuralOperatorNet as SFNO"
    ]
   },
   {

torch_harmonics/__init__.py

@@ -34,3 +34,4 @@ __version__ = '0.6.0'
 from .sht import RealSHT, InverseRealSHT, RealVectorSHT, InverseRealVectorSHT
 from . import quadrature
 from . import random_fields
+import examples

torch_harmonics/examples/__init__.py

+# coding=utf-8
+
+# 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.
+#
+
+from .pde_sphere import SphereSolver
+from .shallow_water_equations import ShallowWaterSolver
\ No newline at end of file

examples/pde_sphere.py → torch_harmonics/examples/pde_sphere.py

@@ -40,8 +40,6 @@ import torch_harmonics as harmonics
 
 import numpy as np
 import matplotlib.pyplot as plt
-from matplotlib import cm
-import matplotlib.tri as mtri
 
 try:
     import cartopy.crs as ccrs
@@ -170,7 +168,7 @@ class SphereSolver(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=25.0)
 
             #ax = plt.gca(projection=proj, frameon=True)
             ax = fig.add_subplot(projection=proj)

torch_harmonics/examples/sfno/__init__.py

+# coding=utf-8
+
+# 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.
+#
+
+from .models.sfno import SphericalFourierNeuralOperatorNet
+from .utils.pde_dataset import PdeDataset

torch_harmonics/examples/sfno/models/__init__.py

+# coding=utf-8
+
+# 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.
+#
\ No newline at end of file

examples/sfno_examples/models/layers.py → torch_harmonics/examples/sfno/models/layers.py

@@ -47,10 +47,10 @@ from typing import Optional
 import math
 
 from torch_harmonics import *
-from models.contractions import *
-from models.activations import *
+from .contractions import *
+from .activations import *
 
-from models.factorizations import get_contract_fun
+from .factorizations import get_contract_fun
 
 # # import FactorizedTensor from tensorly for tensorized operations
 # import tensorly as tl

examples/sfno_examples/models/sfno.py → torch_harmonics/examples/sfno/models/sfno.py

@@ -35,7 +35,7 @@ from apex.normalization import FusedLayerNorm
 
 from torch_harmonics import *
 
-from models.layers import *
+from .layers import *
 
 class SpectralFilterLayer(nn.Module):
     """

examples/shallow_water_equations.py → torch_harmonics/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=25.0)
 
             #ax = plt.gca(projection=proj, frameon=True)
             ax = fig.add_subplot(projection=proj)