adding types hints and a pytest

Author: scaomath

examples/check_SFNO_shapes.py

@@ -0,0 +1,60 @@
+import torch
+from fno.sfno import SFNO
+
+
+if __name__ == "__main__":
+    """
+    testing the arbitrary sizes inference for both
+    spatial and temporal dimensions of SFNO
+    """
+    modes = 8
+    modes_t = 2
+    width = 10
+    bsz = 5
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    sizes = [(n, n, n_t) for (n, n_t) in zip([64, 128, 256], [5, 10, 20])]
+    model = SFNO(modes, modes, modes_t, width, 
+                 latent_steps=3).to(device)
+    x = torch.randn(bsz, *sizes[0]).to(device)
+    _ = model(x)
+
+    try:
+        from torchinfo import summary
+
+        """
+        torchinfo has not resolve the complex number problem
+        """
+        summary(model, input_size=(bsz, *sizes[-1]))
+    except:
+        raise ImportError(
+            "torchinfo is not installed, please install it to get the model summary"
+        )
+    del model
+
+    print("\n" * 3)
+    for k, size in enumerate(sizes):
+        torch.cuda.empty_cache()
+        model = SFNO(modes, modes, modes_t, width, latent_steps=3).to(device)
+        model.add_latent_hook("activations")
+        x = torch.randn(bsz, *size).to(device)
+        pred = model(x)
+        print(f"\n\ninput shape:  {list(x.size())}")
+        print(f"output shape: {list(pred.size())}")
+        for k, v in model.latent_tensors.items():
+            print(k, list(v.shape))
+        del model
+
+    print("\n")
+    # test evaluation speed
+    from time import time
+
+    torch.cuda.empty_cache()
+    model = SFNO(modes, modes, modes_t, width, latent_steps=3).to(device)
+    model.eval()
+    x = torch.randn(bsz, *sizes[1]).to(device)
+    start_time = time()
+    for _ in range(100):
+        pred = model(x)
+    end_time = time()
+    print(f"Average eval for time: {(end_time - start_time) / 100:.6f} seconds")
+    del model

fno/README.md

@@ -38,6 +38,16 @@ Generate the isotropic turbulence in [1] with the inverse cascade frequency sign
 
 ## Training and evaluation scripts
 
+### VSCode workspace for development
+Please add the following setting to your VSCode workspace setting:
+```json
+"settings": {
+		"terminal.integrated.env.osx": {"PYTHONPATH": "${workspaceFolder}"},
+		"terminal.integrated.env.linux": {"PYTHONPATH": "${workspaceFolder}"},
+		"jupyter.notebookFileRoot": "${workspaceFolder}",
+	}
+```
+
 
 ### Testing the arbitrary input and output discretization sizes (including time)
 Run the part below `__name__ == "__main__"` in [`sfno.py`](sfno.py)

fno/base.py

@@ -13,7 +13,7 @@
 from copy import deepcopy
 
 from functools import partial
-from typing import List
+from typing import List, Union, Tuple
 
 import torch
 import torch.fft as fft
@@ -24,6 +24,17 @@
 
 conv_dict = {1: nn.Conv1d, 2: nn.Conv2d, 3: nn.Conv3d}
 
+ACTIVATION_FUNCTIONS = [
+    'CELU', 'ELU', 'GELU', 'GLU', 'Hardtanh', 'Hardshrink', 'Hardsigmoid', 
+    'Hardswish', 'LeakyReLU', 'LogSigmoid', 'MultiheadAttention', 'PReLU', 
+    'ReLU', 'ReLU6', 'RReLU', 'SELU', 'SiLU', 'Sigmoid', 'SoftPlus', 
+    'Softmax', 'Softmax2d', 'Softshrink', 'Softsign', 'Tanh', 'Tanhshrink',
+    'Threshold', 'Mish'
+]
+
+# Type hint for activation functions
+ActivationType = Union[str]
+
 
 class LayerNormnd(nn.GroupNorm):
     """
@@ -50,28 +61,31 @@ def forward(self, v: torch.Tensor):
         return super().forward(v)
 
 
-class MLP(nn.Module):
+class PointwiseFFN(nn.Module):
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        mid_channels,
-        activation: str = "GELU",
+        in_channels: int,
+        out_channels: int,
+        mid_channels: int,
+        activation: ActivationType = "ReLU",
         dim: int = 3,
     ):
-        super(MLP, self).__init__()
+        super().__init__()
+        """
+        Pointwisely-applied 2-layer FFN with a channel expansion
+        """
 
         if dim not in conv_dict:
             raise ValueError(f"Unsupported dimension: {dim}, expected 1, 2, or 3")
 
         Conv = conv_dict[dim]
-        self.mlp1 = Conv(in_channels, mid_channels, 1)
-        self.mlp2 = Conv(mid_channels, out_channels, 1)
+        self.linear1 = Conv(in_channels, mid_channels, 1)
+        self.linear2 = Conv(mid_channels, out_channels, 1)
         self.activation = getattr(nn, activation)()
 
     def forward(self, v: torch.Tensor):
-        for block in [self.mlp1, self.activation, self.mlp2]:
-            v = block(v)
+        for b in [self.linear1, self.activation, self.linear2]:
+            v = b(v)
         return v
 
 
@@ -169,13 +183,13 @@ def forward(self, v, out_mesh_size=None, **kwargs):
         return v
 
 
-class FNO(nn.Module):
+class FNOBase(nn.Module):
     def __init__(
         self,
         *,
         num_spectral_layers: int = 4,
         fft_norm="backward",
-        activation: str = "ReLU",
+        activation: ActivationType = "ReLU",
         spatial_padding: int = 0,
         channel_expansion: int = 4,
         spatial_random_feats: bool = False,
@@ -199,7 +213,7 @@ def __init__(
 
         self.spatial_padding = spatial_padding
         self.fft_norm = fft_norm
-        self.activation_name = activation
+        self.activation = activation
         self.spatial_random_feats = spatial_random_feats
         self.lift_activation = lift_activation
         self.channel_expansion = channel_expansion
@@ -228,10 +242,10 @@ def _set_spectral_layers(
         num_layers: int,
         modes: List[int],
         width: int,
-        activation: str,
-        spectral_conv: nn.Module,
-        mlp: nn.Module,
-        linear: nn.Module,
+        activation: ActivationType,
+        spectral_conv: SpectralConv,
+        mlp: PointwiseFFN,
+        linear: Union[nn.Conv1d, nn.Conv2d, nn.Conv3d],
         channel_expansion: int = 4,
     ) -> None:
         """
@@ -283,8 +297,4 @@ def double(self):
         return self
 
     def forward(self, *args, **kwargs):
-        """
-        if out_steps is None, it will try to use self.out_steps
-        if self.out_steps is None, it will use the temporal dimension of the input
-        """
         raise NotImplementedError("Subclasses of FNO must implement the forward method")