revamping the examples folder structure

examples/train_sfno.py

@@ -44,7 +44,7 @@ import pandas as pd
 
 import matplotlib.pyplot as plt
 
-from torch_harmonics.examples.sfno import PdeDataset
+from torch_harmonics.examples import PdeDataset
 from torch_harmonics import RealSHT
 
 # wandb logging
@@ -409,24 +409,8 @@ def main(train=True, load_checkpoint=False, enable_amp=False, log_grads=0):
     models = {}
     metrics = {}
 
-    from torch_harmonics.examples.sfno import SphericalFourierNeuralOperatorNet as SFNO
-    from torch_harmonics.examples.sfno import LocalSphericalNeuralOperatorNet as LSNO
-
-    # models[f"sfno_sc2_layers4_e32_nomlp"] = partial(
-    #     SFNO,
-    #     img_size=(nlat, nlon),
-    #     grid="equiangular",
-    #     # hard_thresholding_fraction=0.8,
-    #     num_layers=4,
-    #     scale_factor=2,
-    #     embed_dim=32,
-    #     operator_type="driscoll-healy",
-    #     activation_function="gelu",
-    #     big_skip=True,
-    #     pos_embed=False,
-    #     use_mlp=False,
-    #     normalization_layer="none",
-    # )
+    from torch_harmonics.examples.models import SphericalFourierNeuralOperatorNet as SFNO
+    from torch_harmonics.examples.models import LocalSphericalNeuralOperatorNet as LSNO
 
     models[f"sfno_sc2_layers4_e32_nomlp_leggauss"] = partial(
         SFNO,
@@ -444,21 +428,21 @@ def main(train=True, load_checkpoint=False, enable_amp=False, log_grads=0):
         normalization_layer="none",
     )
 
-    # models[f"lsno_sc1_layers4_e32_nomlp"] = partial(
-    #     LSNO,
-    #     spectral_transform="sht",
-    #     img_size=(nlat, nlon),
-    #     grid="equiangular",
-    #     num_layers=4,
-    #     scale_factor=2,
-    #     embed_dim=32,
-    #     operator_type="driscoll-healy",
-    #     activation_function="gelu",
-    #     big_skip=True,
-    #     pos_embed=False,
-    #     use_mlp=False,
-    #     normalization_layer="none",
-    # )
+    models[f"lsno_sc1_layers4_e32_nomlp"] = partial(
+        LSNO,
+        spectral_transform="sht",
+        img_size=(nlat, nlon),
+        grid="equiangular",
+        num_layers=4,
+        scale_factor=2,
+        embed_dim=32,
+        operator_type="driscoll-healy",
+        activation_function="gelu",
+        big_skip=True,
+        pos_embed=False,
+        use_mlp=False,
+        normalization_layer="none",
+    )
 
     # iterate over models and train each model
     root_path = os.path.dirname(__file__)

torch_harmonics/examples/__init__.py

@@ -2,7 +2,7 @@
 
 # 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:
 #
@@ -30,4 +30,6 @@
 #
 
 from .pde_sphere import SphereSolver
-from .shallow_water_equations import ShallowWaterSolver
+from .shallow_water_equations import ShallowWaterSolver
\ No newline at end of file
+
+from .pde_dataset import PdeDataset
\ No newline at end of file

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

@@ -2,7 +2,7 @@
 
 # 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:
 #
@@ -27,4 +27,7 @@
 # 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
+
+from .sfno import SphericalFourierNeuralOperatorNet
+from .lsno import LocalSphericalNeuralOperatorNet

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

@@ -36,7 +36,7 @@ from torch.utils.checkpoint import checkpoint
 import math
 
 from torch_harmonics import *
-from .activations import *
+from ._activations import *
 
 def _no_grad_trunc_normal_(tensor, mean, std, a, b):
     # Cut & paste from PyTorch official master until it's in a few official releases - RW

torch_harmonics/examples/sfno/models/local_sfno.py → torch_harmonics/examples/models/lsno.py

@@ -36,7 +36,7 @@ import torch.amp as amp
 from torch_harmonics import RealSHT, InverseRealSHT
 from torch_harmonics import DiscreteContinuousConvS2, DiscreteContinuousConvTransposeS2
 
-from .layers import *
+from ._layers import *
 
 from functools import partial
 

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 .utils.pde_dataset import PdeDataset
-from .models.sfno import SphericalFourierNeuralOperatorNet
-from .models.local_sfno import LocalSphericalNeuralOperatorNet