data

.gitattributes

@@ -36,3 +36,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 data/class_70b.npy filter=lfs diff=lfs merge=lfs -text
 data/order_70b.npy filter=lfs diff=lfs merge=lfs -text
 data/species_70b.npy filter=lfs diff=lfs merge=lfs -text
+data/pos_embeds_model.npy filter=lfs diff=lfs merge=lfs -text
+model/demo_model.pt filter=lfs diff=lfs merge=lfs -text

app.py

@@ -13,6 +13,8 @@ order_list = list(order[()].keys())
 #genus_list = list(genus[()].keys())
 #family_list = list(family[()].keys())
 
+pos_embed = np.load("data/pos_embed.npy", allow_pickle=True)
+
 def update_fn(val):
     if val=="Class":
         return gr.Dropdown(label="Name", choices=class_list, interactive=True)
@@ -25,6 +27,20 @@ def update_fn(val):
     elif val=="Species":
         return gr.Dropdown(label="Name", choices=species_list, interactive=True)
 
+def pred_fn(taxon, name):
+    if taxon=="Class":
+        text_embeds = clas[()][name]
+    elif taxon=="Order":
+        text_embeds = order[()][name]
+    elif taxon=="Family":
+        text_embeds = family[()][name]
+    elif taxon=="Genus":
+        text_embeds = genus[()][name]
+    elif taxon=="Species":
+        text_embeds = species[()][name]
+
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
     """
@@ -39,5 +55,10 @@ with gr.Blocks() as demo:
     with gr.Row():
         submit_button = gr.Button("Run Model")
     
+    with gr.Row():
+        pred = gr.Image(label="Predicted Heatmap", visible=False)
+    
+    submit_button.click(pred_fn, inputs=[inp, out], outputs=[pred])
     
+
 demo.launch()

data/pos_embeds_model.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e010369a67f1d2946dd787494a65f88c1ed79d1cc6d4a5be3f5ac98568492630
+size 829440128

model/demo_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5d592a3a09086658aeac9b51574f11977962c0f2d5703e0225c3a236be4592d
+size 76024944

requirements.txt

@@ -0,0 +1,4 @@
+numpy==1.23.4
+torch==2.0.1
+rasterio==1.3.8
+einops==0.6.1

sfno_encoder.py

@@ -0,0 +1,479 @@
+import torch
+import torch.nn as nn
+
+from torch_harmonics import *
+
+from torch_harmonics.examples.sfno.models.layers import *
+
+from functools import partial
+
+from einops import repeat
+
+import numpy as np
+
+class SpectralFilterLayer(nn.Module):
+    """
+    Fourier layer. Contains the convolution part of the FNO/SFNO
+    """
+
+    def __init__(
+        self,
+        forward_transform,
+        inverse_transform,
+        embed_dim,
+        filter_type = 'non-linear',
+        operator_type = 'diagonal',
+        sparsity_threshold = 0.0,
+        use_complex_kernels = True,
+        hidden_size_factor = 2,
+        factorization = None,
+        separable = False,
+        rank = 1e-2,
+        complex_activation = 'real',
+        spectral_layers = 1,
+        drop_rate = 0):
+        super(SpectralFilterLayer, self).__init__() 
+
+        if filter_type == 'non-linear' and isinstance(forward_transform, RealSHT):
+            self.filter = SpectralAttentionS2(forward_transform,
+                                              inverse_transform,
+                                              embed_dim,
+                                              operator_type = operator_type,
+                                              sparsity_threshold = sparsity_threshold,
+                                              hidden_size_factor = hidden_size_factor,
+                                              complex_activation = complex_activation,
+                                              spectral_layers = spectral_layers,
+                                              drop_rate = drop_rate,
+                                              bias = False)
+
+        elif filter_type == 'non-linear' and isinstance(forward_transform, RealFFT2):
+            self.filter = SpectralAttention2d(forward_transform,
+                                              inverse_transform,
+                                              embed_dim,
+                                              sparsity_threshold = sparsity_threshold,
+                                              use_complex_kernels = use_complex_kernels,
+                                              hidden_size_factor = hidden_size_factor,
+                                              complex_activation = complex_activation,
+                                              spectral_layers = spectral_layers,
+                                              drop_rate = drop_rate,
+                                              bias = False)
+
+        elif filter_type == 'linear':
+            self.filter = SpectralConvS2(forward_transform,
+                                         inverse_transform,
+                                         embed_dim,
+                                         embed_dim,
+                                         operator_type = operator_type,
+                                         rank = rank,
+                                         factorization = factorization,
+                                         separable = separable,
+                                         bias = True)
+
+        else:
+            raise(NotImplementedError)
+
+    def forward(self, x):
+        return self.filter(x)
+
+class SphericalFourierNeuralOperatorBlock(nn.Module):
+    """
+    Helper module for a single SFNO/FNO block. Can use both FFTs and SHTs to represent either FNO or SFNO blocks.
+    """
+    def __init__(
+            self,
+            forward_transform,
+            inverse_transform,
+            embed_dim,
+            filter_type = 'non-linear',
+            operator_type = 'diagonal',
+            mlp_ratio = 2.,
+            drop_rate = 0.,
+            drop_path = 0.,
+            act_layer = nn.GELU,
+            norm_layer = (nn.LayerNorm, nn.LayerNorm),
+            sparsity_threshold = 0.0,
+            use_complex_kernels = True,
+            factorization = None,
+            separable = False,
+            rank = 128,
+            inner_skip = 'linear',
+            outer_skip = None, # None, nn.linear or nn.Identity
+            concat_skip = False,
+            use_mlp = True,
+            complex_activation = 'real',
+            spectral_layers = 3):
+        super(SphericalFourierNeuralOperatorBlock, self).__init__()
+        
+        # norm layer
+        self.norm0 = norm_layer[0]() #((h,w))
+
+        # convolution layer
+        self.filter = SpectralFilterLayer(forward_transform,
+                                          inverse_transform,
+                                          embed_dim,
+                                          filter_type,
+                                          operator_type = operator_type,
+                                          sparsity_threshold = sparsity_threshold,
+                                          use_complex_kernels = use_complex_kernels,
+                                          hidden_size_factor = mlp_ratio,
+                                          factorization = factorization,
+                                          separable = separable,
+                                          rank = rank,
+                                          complex_activation = complex_activation,
+                                          spectral_layers = spectral_layers,
+                                          drop_rate = drop_rate)
+
+        if inner_skip == 'linear':
+            self.inner_skip = nn.Conv2d(embed_dim, embed_dim, 1, 1)
+        elif inner_skip == 'identity':
+            self.inner_skip = nn.Identity()
+
+        self.concat_skip = concat_skip
+
+        if concat_skip and inner_skip is not None:
+            self.inner_skip_conv = nn.Conv2d(2*embed_dim, embed_dim, 1, bias=False)
+
+        if filter_type == 'linear' or filter_type == 'local':
+            self.act_layer = act_layer()
+        
+        # dropout
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        # norm layer
+        self.norm1 = norm_layer[1]() #((h,w))
+        
+        if use_mlp == True:
+            mlp_hidden_dim = int(embed_dim * mlp_ratio)
+            self.mlp = MLP(in_features = embed_dim,
+                           hidden_features = mlp_hidden_dim,
+                           act_layer = act_layer,
+                           drop_rate = drop_rate,
+                           checkpointing = False)
+
+        if outer_skip == 'linear':
+            self.outer_skip = nn.Conv2d(embed_dim, embed_dim, 1, 1)
+        elif outer_skip == 'identity':
+            self.outer_skip = nn.Identity()
+
+        if concat_skip and outer_skip is not None:
+            self.outer_skip_conv = nn.Conv2d(2*embed_dim, embed_dim, 1, bias=False)
+
+    def forward(self, x):
+        
+        x = self.norm0(x)
+
+        x, residual = self.filter(x)
+
+        if hasattr(self, 'inner_skip'):
+            if self.concat_skip:
+                x = torch.cat((x, self.inner_skip(residual)), dim=1)
+                x = self.inner_skip_conv(x)
+            else:
+                x = x + self.inner_skip(residual)
+
+        if hasattr(self, 'act_layer'):
+            x = self.act_layer(x)
+
+        x = self.norm1(x)
+
+        if hasattr(self, 'mlp'):
+            x = self.mlp(x)
+
+        x = self.drop_path(x)
+
+        if hasattr(self, 'outer_skip'):
+            if self.concat_skip:
+                x = torch.cat((x, self.outer_skip(residual)), dim=1)
+                x = self.outer_skip_conv(x)
+            else:
+                x = x + self.outer_skip(residual)
+
+        return x
+
+class SphericalFourierNeuralOperatorNet(nn.Module):
+    """
+    SphericalFourierNeuralOperator module. Can use both FFTs and SHTs to represent either FNO or SFNO,
+    both linear and non-linear variants.
+
+    Parameters
+    ----------
+    filter_type : str, optional
+        Type of filter to use ('linear', 'non-linear'), by default "linear"
+    spectral_transform : str, optional
+        Type of spectral transformation to use, by default "sht"
+    operator_type : str, optional
+        Type of operator to use ('vector', 'diagonal'), by default "vector"
+    img_shape : tuple, optional
+        Shape of the input channels, by default (128, 256)
+    scale_factor : int, optional
+        Scale factor to use, by default 3
+    in_chans : int, optional
+        Number of input channels, by default 3
+    out_chans : int, optional
+        Number of output channels, by default 3
+    embed_dim : int, optional
+        Dimension of the embeddings, by default 256
+    num_layers : int, optional
+        Number of layers in the network, by default 4
+    activation_function : str, optional
+        Activation function to use, by default "gelu"
+    encoder_layers : int, optional
+        Number of layers in the encoder, by default 1
+    use_mlp : int, optional
+        Whether to use MLP, by default True
+    mlp_ratio : int, optional
+        Ratio of MLP to use, by default 2.0
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    drop_path_rate : float, optional
+        Dropout path rate, by default 0.0
+    sparsity_threshold : float, optional
+        Threshold for sparsity, by default 0.0
+    normalization_layer : str, optional
+        Type of normalization layer to use ("layer_norm", "instance_norm", "none"), by default "instance_norm"
+    hard_thresholding_fraction : float, optional
+        Fraction of hard thresholding (frequency cutoff) to apply, by default 1.0
+    use_complex_kernels : bool, optional
+        Whether to use complex kernels, by default True
+    big_skip : bool, optional
+        Whether to add a single large skip connection, by default True
+    rank : float, optional
+        Rank of the approximation, by default 1.0
+    factorization : Any, optional
+        Type of factorization to use, by default None
+    separable : bool, optional
+        Whether to use separable convolutions, by default False
+    rank : (int, Tuple[int]), optional
+        If a factorization is used, which rank to use. Argument is passed to tensorly
+    complex_activation : str, optional
+        Type of complex activation function to use, by default "real"
+    spectral_layers : int, optional
+        Number of spectral layers, by default 3
+    pos_embed : bool, optional
+        Whether to use positional embedding, by default True
+
+    Example:
+    --------
+    >>> model = SphericalFourierNeuralOperatorNet(
+    ...         img_shape=(128, 256),
+    ...         scale_factor=4,
+    ...         in_chans=2,
+    ...         out_chans=2,
+    ...         embed_dim=16,
+    ...         num_layers=2,
+    ...         encoder_layers=1,
+    ...         num_blocks=4,
+    ...         spectral_layers=2,
+    ...         use_mlp=True,)
+    >>> model(torch.randn(1, 2, 128, 256)).shape
+    torch.Size([1, 2, 128, 256])
+    """
+
+    def __init__(
+            self,
+            filter_type = 'linear',
+            spectral_transform = 'sht',
+            operator_type = 'vector',
+            img_size = (128, 256),
+            scale_factor = 4,
+            in_chans = 3,
+            out_chans = 3,
+            embed_dim = 256,
+            num_layers = 4,
+            activation_function = 'gelu',
+            encoder_layers = 1,
+            use_mlp = True,
+            mlp_ratio = 2.,
+            drop_rate = 0.,
+            drop_path_rate = 0.,
+            sparsity_threshold = 0.0,
+            normalization_layer = 'instance_norm',
+            hard_thresholding_fraction = 1.0,
+            use_complex_kernels = True,
+            big_skip = False,
+            factorization = None,
+            separable = False,
+            rank = 128,
+            complex_activation = 'real',
+            spectral_layers = 2,
+            pos_embed = True
+            ):
+
+        super(SphericalFourierNeuralOperatorNet, self).__init__()
+
+        self.filter_type = filter_type
+        self.spectral_transform = spectral_transform
+        self.operator_type = operator_type
+        self.img_size = img_size
+        self.scale_factor = scale_factor
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.embed_dim = self.num_features = embed_dim
+        self.pos_embed_dim = self.embed_dim
+        self.num_layers = num_layers
+        self.hard_thresholding_fraction = hard_thresholding_fraction
+        self.normalization_layer = normalization_layer
+        self.use_mlp = use_mlp
+        self.encoder_layers = encoder_layers
+        self.big_skip = big_skip
+        self.factorization = factorization
+        self.separable = separable,
+        self.rank = rank
+        self.complex_activation = complex_activation
+        self.spectral_layers = spectral_layers
+
+        # activation function
+        if activation_function == 'relu':
+            self.activation_function = nn.ReLU
+        elif activation_function == 'gelu':
+            self.activation_function = nn.GELU
+        else:
+            raise ValueError(f"Unknown activation function {activation_function}")
+
+        # compute downsampled image size
+        self.h = self.img_size[0] // scale_factor
+        self.w = self.img_size[1] // scale_factor
+
+        # dropout
+        self.pos_drop = nn.Dropout(p=drop_rate) if drop_rate > 0. else nn.Identity()
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, self.num_layers)]
+
+        # pick norm layer
+        if self.normalization_layer == "layer_norm":
+            norm_layer0 = partial(nn.LayerNorm, normalized_shape=(self.img_size[0], self.img_size[1]), eps=1e-6)
+            norm_layer1 = partial(nn.LayerNorm, normalized_shape=(self.h, self.w), eps=1e-6)
+        elif self.normalization_layer == "instance_norm":
+            norm_layer0 = partial(nn.InstanceNorm2d, num_features=self.embed_dim, eps=1e-6, affine=True, track_running_stats=False)
+            norm_layer1 = norm_layer0
+        elif self.normalization_layer == "none":
+            norm_layer0 = nn.Identity
+            norm_layer1 = norm_layer0
+        else:
+            raise NotImplementedError(f"Error, normalization {self.normalization_layer} not implemented.") 
+
+        if pos_embed:
+            self.pos_embed = nn.Parameter(torch.zeros(1, self.embed_dim, self.img_size[0], self.img_size[1]))
+            #self.pos_embed = posemb_sincos_2d(900, 1800, 128)
+            pass
+            #x = torch.linspace(-np.pi, np.pi, 900)
+            #y = torch.linspace(-np.pi, np.pi, 1800)
+            #x, y = torch.meshgrid(x, y)
+            #self.pos_embed = torch.stack((torch.sin(x), torch.sin(y), torch.cos(x), torch.cos(y)), dim=0).unsqueeze(0).cuda()
+            #self.pos_embed = nn.Parameter(torch.zeros(1, self.embed_dim, self.img_size[0], self.img_size[1]))
+            #self.pos_direct = nn.Conv2d(4, self.embed_dim, 1, bias=False)
+        else:
+            self.pos_embed = None
+
+        # encoder
+        """encoder_hidden_dim = self.embed_dim
+        current_dim = self.in_chans
+        encoder_modules = []
+        for i in range(self.encoder_layers):
+            encoder_modules.append(nn.Conv2d(current_dim, encoder_hidden_dim, 1, bias=True))
+            encoder_modules.append(self.activation_function())
+            current_dim = encoder_hidden_dim
+        encoder_modules.append(nn.Conv2d(current_dim, self.embed_dim, 1, bias=False))
+        self.encoder = nn.Sequential(*encoder_modules)"""
+        
+        # prepare the spectral transform
+        if self.spectral_transform == 'sht':
+
+            modes_lat = int(self.h * self.hard_thresholding_fraction)
+            modes_lon = int((self.w // 2 + 1) * self.hard_thresholding_fraction)
+
+            self.trans_down = RealSHT(*self.img_size, lmax=modes_lat, mmax=modes_lon, grid='equiangular').float()
+            self.itrans_up  = InverseRealSHT(*self.img_size, lmax=modes_lat, mmax=modes_lon, grid='equiangular').float()
+            self.trans      = RealSHT(self.h, self.w, lmax=modes_lat, mmax=modes_lon, grid='legendre-gauss').float()
+            self.itrans     = InverseRealSHT(self.h, self.w, lmax=modes_lat, mmax=modes_lon, grid='legendre-gauss').float()
+
+        elif self.spectral_transform == 'fft':
+
+            modes_lat = int(self.h * self.hard_thresholding_fraction)
+            modes_lon = int((self.w // 2 + 1) * self.hard_thresholding_fraction)
+
+            self.trans_down = RealFFT2(*self.img_size, lmax=modes_lat, mmax=modes_lon).float()
+            self.itrans_up  = InverseRealFFT2(*self.img_size, lmax=modes_lat, mmax=modes_lon).float()
+            self.trans      = RealFFT2(self.h, self.w, lmax=modes_lat, mmax=modes_lon).float()
+            self.itrans     = InverseRealFFT2(self.h, self.w, lmax=modes_lat, mmax=modes_lon).float()
+            
+        else:
+            raise(ValueError('Unknown spectral transform'))
+
+        self.blocks = nn.ModuleList([])
+        for i in range(self.num_layers):
+
+            first_layer = i == 0
+            last_layer = i == self.num_layers-1
+
+            forward_transform = self.trans_down if first_layer else self.trans
+            inverse_transform = self.itrans_up if last_layer else self.itrans
+
+            inner_skip = 'linear'
+            outer_skip = 'identity'
+
+            if first_layer:
+                norm_layer = (norm_layer0, norm_layer1)
+            elif last_layer:
+                norm_layer = (norm_layer1, norm_layer0)
+            else:
+                norm_layer = (norm_layer1, norm_layer1)
+
+            block = SphericalFourierNeuralOperatorBlock(forward_transform,
+                                                        inverse_transform,
+                                                        self.embed_dim,
+                                                        filter_type = filter_type,
+                                                        operator_type = self.operator_type,
+                                                        mlp_ratio = mlp_ratio,
+                                                        drop_rate = drop_rate,
+                                                        drop_path = dpr[i],
+                                                        act_layer = self.activation_function,
+                                                        norm_layer = norm_layer,
+                                                        sparsity_threshold = sparsity_threshold,
+                                                        use_complex_kernels = use_complex_kernels,
+                                                        inner_skip = inner_skip,
+                                                        outer_skip = outer_skip,
+                                                        use_mlp = use_mlp,
+                                                        factorization = self.factorization,
+                                                        separable = self.separable,
+                                                        rank = self.rank,
+                                                        complex_activation = self.complex_activation,
+                                                        spectral_layers = self.spectral_layers)
+
+            self.blocks.append(block)
+
+        # trunc_normal_(self.pos_embed, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
+            trunc_normal_(m.weight, std=.02)
+            #nn.init.normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def forward_features(self, x):
+
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+            
+        return x
+
+    def forward(self, x):
+
+        #if self.big_skip:
+            #residual = x
+
+        #x = self.encoder(x)
+
+        #x = x + self.pos_embed 
+        x = self.pos_embed
+        
+        x = self.forward_features(x)
+
+        return x