Fix geo embeding

climanet/predict.py

@@ -44,8 +44,8 @@ def _save_netcdf(
     full_predictions = np.full(
         (len(times), len(lats), len(lons)), np.nan, dtype=predictions.dtype
     )
-    for i, patch_idx in enumerate(indices):
-        month_start, lat_start, lon_start = base_dataset.patch_indices[patch_idx]
+    for i, data_idx in enumerate(indices):
+        month_start, lat_start, lon_start = base_dataset.crop_indices[data_idx]
         full_predictions[
             month_start : month_start + M,
             lat_start : lat_start + H,
@@ -82,17 +82,17 @@ def _move_batch_to_device(batch: dict, device: str):
 def _run_one_batch(model: torch.nn.Module, batch: dict, device: str):
     batch = _move_batch_to_device(batch, device)
     pred = model(
-        batch["daily_patch"],
-        batch["daily_mask_patch"],
-        batch["daily_timef_patch"],
-        batch["land_mask_patch"],
-        batch["geo_pos_embedding_patch"],
-        batch["scale_feature_patch"],
+        batch["input_data"],
+        batch["input_data_mask"],
+        batch["input_data_timef"],
+        batch["land_mask"],
+        batch["geo_pos_embedding"],
+        batch["scale_feature"],
         batch["padded_days_mask"],
     )  # (B, M, H, W)
 
     # Compute masked loss
-    loss = compute_masked_loss(pred, batch["monthly_patch"], batch["land_mask_patch"])
+    loss = compute_masked_loss(pred, batch["monthly_data"], batch["land_mask"])
     return loss, pred
 
 
@@ -143,7 +143,7 @@ def predict_monthly_var(
     # Initialize an empty list to store predictions
     base_dataset = dataset.dataset if hasattr(dataset, "dataset") else dataset
 
-    M, H, W = base_dataset.patch_size
+    M, H, W = base_dataset.crop_size
     all_predictions = torch.empty(len(dataset), M, H, W, device=device)
 
     # Set up logging

climanet/st_encoder_decoder.py

@@ -41,9 +41,6 @@ def __init__(self, in_chans=1, embed_dim=128, patch_size=(1, 4, 4)):
             2 * in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
         )
 
-        # norm is LayerNorm over the embedding dimension to normalize patch embeddings
-        self.norm = nn.LayerNorm(embed_dim)
-
     def forward(self, x, mask):
         """Forward pass with masking support via an additional validity channel.
         Args:
@@ -70,7 +67,6 @@ def forward(self, x, mask):
         x = self.proj(x)
         x = x.flatten(2).transpose(1, 2)
 
-        x = self.norm(x)
         return x
 
 
@@ -267,11 +263,11 @@ def __init__(self, embed_dim=128, dropout=0.0, chunk_size=1024):
 
         # Day scorer (within each month)
         self.day_scorer = nn.Sequential(
-            nn.LayerNorm(embed_dim),  # normalizing features
-            nn.Linear(embed_dim, embed_dim),  # learns temporal feature transformation
-            nn.GELU(),  # adds non-linearity to capture complex temporal patterns
+            nn.LayerNorm(2 * embed_dim),
+            nn.Linear(2 * embed_dim, embed_dim),
+            nn.GELU(),
             nn.Dropout(dropout),
-            nn.Linear(embed_dim, 1),  # project to a single score
+            nn.Linear(embed_dim, 1),
         )
 
         # Cross month mixing
@@ -330,13 +326,19 @@ def forward(self, x, M, time_features, padded_days_mask=None):
         month_emb = self.month_embed(time_features)
         token_emb = temp_emb + month_emb
 
-        day_logits = self.day_scorer(token_emb).squeeze(-1)
+        # Make day attention depend on both the actual daily data
+        # and the temporal/calendar information.
+        token_emb_seq = token_emb.unsqueeze(1)  # (B, 1, M, Tp, C)
+        day_score_input = torch.cat(
+            [seq,token_emb_seq.expand(-1, HW, -1, -1, -1)], dim=-1,
+        )
+
+        day_logits = self.day_scorer(day_score_input).squeeze(-1)
 
         if padded_days_mask is not None:
-            day_logits = day_logits.masked_fill(padded_days_mask, float("-inf"))
+            day_logits = day_logits.masked_fill(padded_days_mask.unsqueeze(1), float("-inf"))
 
-        day_w = torch.softmax(day_logits, dim=-1)
-        day_w = day_w.unsqueeze(1).unsqueeze(-1)
+        day_w = torch.softmax(day_logits, dim=-1).unsqueeze(-1)
 
         # avoid implicit expand copies
         seq = seq.reshape(B, HW, M, Tp, C)
@@ -346,7 +348,6 @@ def forward(self, x, M, time_features, padded_days_mask=None):
 
         # avoid broadcast materialization
         aggregated_month_embed = (token_emb_seq * day_w).sum(dim=3)
-
         month_tokens = month_tokens + aggregated_month_embed
 
         z = month_tokens.reshape(B * HW, M, C)
@@ -379,7 +380,6 @@ def __init__(
         patch_h=4,
         patch_w=4,
         hidden=128,
-        overlap=1,
         dropout=0.0,
     ):
         """
@@ -391,14 +391,11 @@ def __init__(
             patch_w: Patch width
             hidden: Hidden dimension in the decoder for mixing channel features.
                 The default is 128, which can be tuned.
-            overlap: Overlap size for deconvolution. It creates smooth blending
-                between adjacent upsampled patches. Default is 1, no overlap at edges.
             dropout: Dropout rate for regularization in the refinement block. Default is 0.0.
         """
         super().__init__()
         self.patch_h = patch_h
         self.patch_w = patch_w
-        self.overlap = overlap
 
         # Mix channel features on the patch grid (Hp, Wp)
         # Input shape: (B, embed_dim, Hp, Wp) → Output shape: (B, hidden, Hp, Wp)
@@ -408,21 +405,13 @@ def __init__(
         self.proj = nn.Conv2d(in_channels, out_channels, kernel_size=1)
 
         # Upsample to full resolution
-        # With kernel = stride + 2*overlap and padding=overlap,
-        # output size is exact: H = Hp*patch_h, W = Wp*patch_w (no output_padding needed).
-        k_h = patch_h + 2 * overlap
-        k_w = patch_w + 2 * overlap
-        # As spatial size increases, channel count decreases to keep computation
-        # manageable; here  hidden // 2 is a design choice.
         in_channels, out_channels = hidden, hidden // 2
         self.deconv = nn.ConvTranspose2d(
             in_channels,
             out_channels,
-            kernel_size=(k_h, k_w),
+            kernel_size=(patch_h, patch_w),
             stride=(patch_h, patch_w),
-            padding=overlap,
-            output_padding=0,
-            bias=True,
+            padding=0,
         )
 
         # Final conv head to get single channel output kernel_size=3 is the most
@@ -474,7 +463,9 @@ def forward(self, latent, M, out_H, out_W, land_mask=None):
         out = self.deconv(out)  # (B*M, hidden//2, H, W)
 
         # Refinement CNN to smooth boundaries
+        refine_input = out
         out = self.refine(out)  # (B*M, hidden//2, H, W)
+        out = out + refine_input  # Residual connection to preserve original features
 
         # Apply final conv head to get single channel output
         out = self.head(out)  # (B*M, 1, H, W)
@@ -568,56 +559,6 @@ def forward(
         return geo_emb
 
 
-class SpatialTransformer(nn.Module):
-    """Spatial Transformer for spatial feature mixing.
-
-    This module applies a standard Transformer encoder to a sequence of spatial tokens
-    (patch embeddings), allowing information to be mixed across all spatial locations.
-
-    Key points:
-        - Uses multi-head self-attention and feedforward layers.
-        - Designed to operate on flattened spatial tokens.
-    """
-
-    def __init__(self, embed_dim=128, depth=2, num_heads=4, mlp_ratio=4.0, dropout=0.0):
-        """Initialize the spatial transformer.
-        Args:
-            embed_dim: Dimension of the embedding. Default is 128.
-                The embedding dimensions are multiples of 64 (e.g., 64, 128,
-                256). This can be tuned.
-            depth: Number of transformer encoder layers. Default is 2. This can be
-                increased for more complex spatial mixing.
-            num_heads: Number of attention heads in each layer. Default is 4.
-                When embed_dim is 128, 4 heads is a common choice.
-            mlp_ratio: Ratio of feedforward hidden dimension to embed_dim. Default is 4.0.
-            dropout: Dropout rate applied to attention and feedforward layers. Default is 0.0.
-        """
-        super().__init__()
-
-        # a single Transformer encoder block that
-        # performs self-attention and feedforward processing
-        encoder_layer = nn.TransformerEncoderLayer(
-            d_model=embed_dim,
-            nhead=num_heads,
-            dim_feedforward=int(embed_dim * mlp_ratio),
-            batch_first=True,
-            dropout=dropout,
-            activation="gelu",
-        )
-        # stack multiple layers to form the full encoder
-        self.enc = nn.TransformerEncoder(encoder_layer, num_layers=depth)
-
-    def forward(self, x):
-        """Forward pass of the spatial transformer.
-        Args:
-            x: Input tensor of shape (B, N, C), where N = number of spatial tokens (H'*W') and
-                C = embedding dimension
-        Returns:
-            Tensor of shape (B, N, C) with spatially mixed features across patches
-        """
-        return self.enc(x)
-
-
 class SpatioTemporalModel(nn.Module):
     """Spatio-Temporal Model for Monthly Prediction.
 
@@ -631,7 +572,7 @@ class SpatioTemporalModel(nn.Module):
     The model pipeline:
         1. Encode spatio-temporal patches using VideoEncoder.
         2. Aggregate temporal information for each spatial patch via TemporalAttentionAggregator.
-        3. Add 2D spatial positional encodings and mix spatial features with SpatialTransformer.
+        3. Add 2D spatial positional encodings.
         4. Decode aggregated tokens into a full-resolution 2D map using MonthlyConvDecoder.
 
     Output:
@@ -644,9 +585,6 @@ def __init__(
         embed_dim=128,
         patch_size=(1, 4, 4),
         hidden=256,
-        overlap=1,
-        spatial_depth=2,
-        spatial_heads=4,
         dropout=0.0,
         sh_dim=96,
         scale_dim=10,
@@ -659,11 +597,8 @@ def __init__(
             embed_dim: Dimension of the patch embedding
             patch_size: Tuple of (T, H, W) patch sizes for temporal and spatial patching
             hidden: Hidden dimension used in the decoder
-            overlap: Overlap for deconvolution in the decoder
             max_H: Maximum spatial height for 2D positional encoding
             max_W: Maximum spatial width for 2D positional encoding
-            spatial_depth: Number of layers in the spatial Transformer
-            spatial_heads: Number of attention heads in the spatial Transformer
             dropout: Dropout rate for regularization in various components. Increase it if there is overfitting.
             sh_dim: Dimension of spherical harmonics based pca of geo-position
             scale_dim: Dimension of patch-level patch-scale features
@@ -691,18 +626,12 @@ def __init__(
             scale_dim=scale_dim,
             embed_dim=embed_dim,
         )
-        self.spatial_tr = SpatialTransformer(
-            embed_dim=embed_dim,
-            depth=spatial_depth,
-            num_heads=spatial_heads,
-            dropout=dropout,
-        )
+
         self.decoder = MonthlyConvDecoder(
             embed_dim=embed_dim,
             patch_h=patch_size[1],
             patch_w=patch_size[2],
             hidden=hidden,
-            overlap=overlap,
             dropout=dropout,
         )
         self.patch_size = patch_size
@@ -726,6 +655,8 @@ def forward(
             daily_timef: Tensor of shape (B, M, T, 2) containing the cyclically phase encoded day-of-year
                 and hour-of-day information for the daily data
             land_mask_patch: Boolean tensor of shape (B, H, W) to mask land areas in the output
+            geo_pos_embedding_patch: Tensor of shape (B, Hp*Wp, sh_embed_dim) containing the geo-positional embeddings for each model patch within the dataset patch
+            scale_feature_patch: Tensor of shape (B, Hp*Wp, scale_dim) containing the scale features for each model patch within the dataset patch
             padded_days_mask: Optional boolean tensor of shape (B, M, T) indicating which day tokens are padded
                  (True for padded tokens). Used to mask out padded tokens in temporal attention.
         Returns:
@@ -787,32 +718,27 @@ def forward(
             )  # (B, M, Hp*Wp, embed_dim)
 
         # Step 3: Add geo position and scale encodings
+        B = geo_pos_embedding_patch.shape[0]
+
         if self.use_checkpoint:
-            geo_emb = checkpoint(
+            geo_emb_flat = checkpoint(
                 self.geo_embedding,
-                geo_pos_embedding_patch,
-                scale_feature_patch,
+                geo_pos_embedding_patch.view(B * Hp * Wp, -1),
+                scale_feature_patch.view(B * Hp * Wp, -1),
                 use_reentrant=False,
             )[:, None, None, :]  # (B,1,1,E)
         else:
-            geo_emb = self.geo_embedding(geo_pos_embedding_patch, scale_feature_patch)[
-                :, None, None, :
-            ]  # (B,1,1,E)
+            geo_emb_flat = self.geo_embedding(
+                geo_pos_embedding_patch.view(B * Hp * Wp, -1),
+                scale_feature_patch.view(B * Hp * Wp, -1)
+            )  # (B*Hp*Wp, E)
 
         # Broadcasting: same geo embedding for all M months at each Hp*Wp location
         # we use weighted mean patch embedding, see `geo_embedding_utils.py`
+        geo_emb = geo_emb_flat.view(B, 1, Hp * Wp, embed_dim)  # (B, 1, Hp*Wp, E)
         x = agg_latent + geo_emb  # (B, M, Hp*Wp, E)
 
-        # Step 4: Spatial mixing with Transformer
-        # spatial transformer input shape = (B, N, C), output shape = (B, N, C) C: embedding dimension
-        # M is folded in B.
-
-        C = x.shape[-1]
-        x = x.reshape(B * M, Hp * Wp, C)
-        x = self.spatial_tr(x)
-        x = x.view(B, M, Hp * Wp, C)
-
-        # Step 5: Decode to full-resolution 2D map
+        # Step 4: Decode to full-resolution 2D map
         # decoder input shape is (B, M*Hp*Wp, C), C: embedding dimension
         # decoder output shape is (B, M, H, W)
         if self.use_checkpoint:

climanet/train.py

@@ -41,17 +41,17 @@ def _move_batch_to_device(batch: dict, device: str):
 def _run_one_batch(model: torch.nn.Module, batch: dict, device):
     batch = _move_batch_to_device(batch, device)
     pred = model(
-        batch["daily_patch"],
-        batch["daily_mask_patch"],
-        batch["daily_timef_patch"],
-        batch["land_mask_patch"],
-        batch["geo_pos_embedding_patch"],
-        batch["scale_feature_patch"],
+        batch["input_data"],
+        batch["input_data_mask"],
+        batch["input_data_timef"],
+        batch["land_mask"],
+        batch["geo_pos_embedding"],
+        batch["scale_feature"],
         batch["padded_days_mask"],
     )  # (B, M, H, W)
 
     # Compute masked loss
-    return compute_masked_loss(pred, batch["monthly_patch"], batch["land_mask_patch"])
+    return compute_masked_loss(pred, batch["monthly_data"], batch["land_mask"])
 
 
 def _load_checkpoint(model, optimizer, loaded_checkpoint):

climanet/tune.py

@@ -26,8 +26,9 @@ def _tune_data_preparation(data_config):
         padded_days_mask=padded_days_mask,
         time_features=time_features,
         land_mask=ray.get(data_config["land_mask_data"]),
-        patch_size=data_config["patch_size"],  # based on the patch_size in model
+        crop_size=data_config["crop_size"],  # based on the patch_size in model
         stride=data_config["stride"],
+        model_patch_size=data_config["model_patch_size"],
         sh_embed_dim=96,
         sh_order_L=10,
         verbose=False,
@@ -73,20 +74,14 @@ def _train(tune_config, static_args):
     set_seed()
 
     patch_size = tune_config["patch_size"]
-    overlap = tune_config["overlap"]
     embed_dim = tune_config["embed_dim"]
     dropout = tune_config["dropout"]
     hidden = tune_config["hidden"]
-    spatial_depth = tune_config["spatial_depth"]
-    spatial_heads = tune_config["spatial_heads"]
     model = SpatioTemporalModel(
         patch_size=(1, patch_size, patch_size),
-        overlap=overlap,
         embed_dim=embed_dim,
         dropout=dropout,
         hidden=hidden,
-        spatial_depth=spatial_depth,
-        spatial_heads=spatial_heads,
     )
 
     _ = train_monthly_model(