68 replace month pe

climanet/dataset.py

@@ -96,10 +96,18 @@ def __init__(
             )
 
         # Convert to tensor once — all __getitem__ calls use these
-        self.daily_t = torch.from_numpy(daily_mt.values.astype(np.float32))  # (M, T=31, H, W)
-        self.monthly_t = torch.from_numpy(monthly_m.values.astype(np.float32))  # (M, H, W)
-        self.padded_days_t = torch.from_numpy(padded_days_mask.values.copy()).bool()  # (M, T=31)
-        self.daily_timef_t = torch.from_numpy(daily_timef.values.astype(np.float32))  # (M, T=31, 4)
+        self.daily_t = torch.from_numpy(
+            daily_mt.values.astype(np.float32)
+        )  # (M, T=31, H, W)
+        self.monthly_t = torch.from_numpy(
+            monthly_m.values.astype(np.float32)
+        )  # (M, H, W)
+        self.padded_days_t = torch.from_numpy(
+            padded_days_mask.values.copy()
+        ).bool()  # (M, T=31)
+        self.daily_timef_t = torch.from_numpy(
+            daily_timef.values.astype(np.float32)
+        )  # (M, T=31, 3)
 
         # Store coordinate arrays
         self.lat_coords = input_da[spatial_dims[0]].to_numpy().copy()
@@ -180,7 +188,9 @@ def _compute_patch_indices(self, M: int, H: int, W: int) -> list:
 
         # Compute patch start indices using stride
         # Ensure we don't go out of bounds
-        m_starts = list(range(0, M - pm + 1, pm))  # Temporal patches are non-overlapping
+        m_starts = list(
+            range(0, M - pm + 1, pm)
+        )  # Temporal patches are non-overlapping
         i_starts = list(range(0, H - ph + 1, sh))
         j_starts = list(range(0, W - pw + 1, sw))
 
@@ -266,7 +276,9 @@ def __getitem__(self, idx):
         monthly_t_patch = self.monthly_t[m : m + pm, i : i + ph, j : j + pw]
 
         # (M, T, H, W) -> (M, T, pH, pW)
-        daily_nan_mask_t_patch = self.daily_nan_mask_t[m : m + pm, :, i : i + ph, j : j + pw].unsqueeze(0)
+        daily_nan_mask_t_patch = self.daily_nan_mask_t[
+            m : m + pm, :, i : i + ph, j : j + pw
+        ].unsqueeze(0)
 
         if self.land_mask_t is not None:
             land_t_patch = self.land_mask_t[i : i + ph, j : j + pw]  # (H, W)
@@ -295,8 +307,10 @@ def __getitem__(self, idx):
             "monthly_patch": monthly_t_patch,  # (pm, pH, pW)
             "daily_mask_patch": daily_mask_t_patch,  # (C=1, pm, T=31, pH, pW)
             "land_mask_patch": land_t_patch,  # (pH,pW) True=Land
-            "daily_timef_patch": self.daily_timef_t[m : m + pm],  # (pm, T=31, 2)
-            "padded_days_mask": self.padded_days_t[m : m + pm],  # (pm, T=31) True=padded
+            "daily_timef_patch": self.daily_timef_t[m : m + pm],  # (pm, T=31, 3)
+            "padded_days_mask": self.padded_days_t[
+                m : m + pm
+            ],  # (pm, T=31) True=padded
             "scale_feature_patch": scale_feature_t,  # (10,)
             "geo_pos_embedding_patch": geo_pos_embedding_t,  # (sh_embed_dim,)
             "sh_embed_dim": self.sh_embed_dim_t,

climanet/predict.py

@@ -23,9 +23,11 @@ def _save_netcdf(predictions: np.ndarray, dataset: Dataset, save_dir: str):
     )
     for i, patch_idx in enumerate(indices):
         month_start, lat_start, lon_start = base_dataset.patch_indices[patch_idx]
-        full_predictions[month_start : month_start + M, lat_start : lat_start + H, lon_start : lon_start + W] = (
-            predictions[i]
-        )
+        full_predictions[
+            month_start : month_start + M,
+            lat_start : lat_start + H,
+            lon_start : lon_start + W,
+        ] = predictions[i]
 
     data_vars = {
         "predictions": (("time", "lat", "lon"), full_predictions),

climanet/st_encoder_decoder.py

@@ -3,7 +3,6 @@
 The main model class is SpatioTemporalModel.
 """
 
-import math
 import torch
 import torch.nn as nn
 from torch.utils.checkpoint import checkpoint
@@ -75,6 +74,64 @@ def forward(self, x, mask):
         return x
 
 
+class CyclicMonthEmbedding(nn.Module):
+    """Cyclical encoding of month using month-of-year as phse.
+
+    This module uses a Fourier base corresponding to the C12 cyclical
+    symmetry group (month based phase representaton.)
+    """
+
+    def __init__(self, embed_dim=128, n_harmonics=6):
+        """
+        Initialize monthly encoding
+
+        Args:
+            embed_dim: Dimension of the embedding.The default is 128.
+                Many vision transformers use embedding dimensions that are multiples
+                of 64 (e.g., 64, 128, 256). This can be tuned.
+            n_harmonics: number of harmonics to consider for the fourier basis.
+                Can be modified, but the deafult value of 6 (sin/cos pairs, so
+                12 functions) represents a complete basis set for the C12 group
+                which the months form.
+        """
+
+        super().__init__()
+
+        self.n_harmonics = n_harmonics
+
+        # set fixed frequencies
+        freq = torch.linspace(1.0, n_harmonics, n_harmonics)
+        self.register_buffer("freq", freq)
+
+        # learned projection into embedding space.
+        # The fourier basis is considerably smaller than
+        # the embedding dimension. This make a learned projection
+        # desirable. No additive bias is allowed
+        self.proj = nn.Linear(2 * n_harmonics, embed_dim, bias=False)
+
+    def forward(self, time_features):
+        """
+        Create month embedding for tokens
+        """
+        # extract moy phase
+        moy_phase = time_features[..., 0]  # (B, M, T)
+
+        x = moy_phase.unsqueeze(-1)  # (B, M, T, 1)
+
+        # (1,1,1,F)
+        freq = self.freq.view(1, 1, 1, -1)
+
+        # apply frequencies
+        x = x * freq
+
+        sinx = torch.sin(x)
+        cosx = torch.cos(x)
+
+        fourier_emb = torch.cat([sinx, cosx], dim=-1)  # (B, M, T, 12)
+
+        return self.proj(fourier_emb)  # (B, M, T, embed_dim)
+
+
 class CyclicTimeEmbedding(nn.Module):
     """Cyclical Temporal encoding using day-of-year and hour-of-day values in
     combination sine and cosine functions
@@ -138,8 +195,8 @@ def forward(self, time_features):
         B, M, T, D = time_features.shape
 
         # extract individual phases from features
-        phase_doy = time_features[..., 0]
-        phase_hod = time_features[..., 1]
+        phase_doy = time_features[..., 1]
+        phase_hod = time_features[..., 2]
         phases = [phase_doy, phase_hod]
 
         # construct cross terms
@@ -169,47 +226,6 @@ def forward(self, time_features):
         return emb_encode
 
 
-class TemporalPositionalEncoding(nn.Module):
-    """Temporal Positional Encoding using sine and cosine functions.
-
-    This module generates fixed (non-learnable) sinusoidal positional encodings
-    for the temporal dimension, following the formulation in
-    "Attention Is All You Need" (Vaswani et al., 2017).
-
-    The returned positional encodings are intended to be added to temporal
-    embeddings by the caller, but this module itself does not perform the addition.
-    """
-
-    def __init__(self, embed_dim=128, max_len=31):
-        """Initialize the temporal positional encoding.
-        Args:
-            embed_dim: Dimension of the embedding.The default is 128.
-                Many vision transformers use embedding dimensions that are multiples
-                of 64 (e.g., 64, 128, 256). This can be tuned.
-            max_len: Maximum length of the temporal dimension to precompute
-            encodings for. Default is 31, which is sufficient for a month of
-            daily data.
-        """
-        super().__init__()
-        pe = torch.zeros(max_len, embed_dim)
-        position = torch.arange(0, max_len).unsqueeze(1)
-        div_term = torch.exp(
-            torch.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim)
-        )
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        self.register_buffer("pe", pe)  # (max_len, embeddim)
-
-    def forward(self, T):
-        """Return positional encodings for a temporal sequence.
-        Args:
-            T: Temporal length (must be <= max_len)
-        Returns:
-            Tensor of shape (T, embed_dim) containing sinusoidal positional encodings
-        """
-        return self.pe[:T]  # (T, embed_dim)
-
-
 class TemporalAttentionAggregator(nn.Module):
     """Temporal attention-based aggregator.
 
@@ -228,24 +244,22 @@ class TemporalAttentionAggregator(nn.Module):
     months.
     """
 
-    def __init__(self, embed_dim=128, max_months=12, dropout=0.0):
+    def __init__(self, embed_dim=128, dropout=0.0):
         """Initialize the temporal attention aggregator.
 
         Args:
             embed_dim: Dimension of the embedding. The default is 128.
                 Many vision transformers use embedding dimensions that are multiples
                 of 64 (e.g., 64, 128, 256). This can be tuned.
-            max_months: Maximum number of months (temporal patches) to precompute
-            encodings for. Default is 12, which is sufficient for a year of monthly data.
             dropout: Dropout rate for regularization in the day scorer and
             cross-month mixing. Default is 0.0. Increase it if there is overfitting.
         """
         super().__init__()
 
         self.time_embed = CyclicTimeEmbedding(embed_dim=embed_dim)
 
-        # Positional encodings for days and months
-        self.pos_months = TemporalPositionalEncoding(embed_dim, max_len=max_months)
+        # cyclical embedding for months
+        self.month_embed = CyclicMonthEmbedding(embed_dim=embed_dim)
 
         # Day scorer (within each month)
         self.day_scorer = nn.Sequential(
@@ -274,10 +288,6 @@ def __init__(self, embed_dim=128, max_months=12, dropout=0.0):
             nn.Linear(4 * embed_dim, embed_dim),
         )
 
-        # Pre-compute and register as buffer — auto-moves with .to(device/dtype)
-        pe = self.pos_months(max_months)  # (max_months, C)
-        self.register_buffer("pe_months_cache", pe)  # tracks device/dtype automatically
-
     def forward(self, x, M, time_features, padded_days_mask=None):
         """
         Args:
@@ -286,7 +296,7 @@ def forward(self, x, M, time_features, padded_days_mask=None):
             T: number of temporal tokens per month after temporal patching (Tp)
             H: spatial height after spatial patching
             W: spatial width after spatial patching
-            time_features: (B,M,T,2) containing cyclically phase encoded DOY and HOD
+            time_features: (B,M,T,3) containing cyclically phase encoded MOY, DOY and HOD
             padded_days_mask: Optional boolean tensor of shape (B, M, T), bool,
                 True indicating which day tokens are padded (because some months
                 have fewer days). This is used to mask out padded tokens in attention computation.
@@ -301,8 +311,8 @@ def forward(self, x, M, time_features, padded_days_mask=None):
         seq = x.reshape(B, M, Tp, HW, C).permute(0, 3, 1, 2, 4)
 
         temp_emb = self.time_embed(time_features)
-        pe_months = self.pe_months_cache[:M]
-        token_emb = temp_emb + pe_months[None, :, None, :]
+        month_emb = self.month_embed(time_features)
+        token_emb = temp_emb + month_emb
 
         day_logits = self.day_scorer(token_emb).squeeze(-1)
 
@@ -319,9 +329,9 @@ def forward(self, x, M, time_features, padded_days_mask=None):
         month_tokens = (seq * day_w).sum(dim=3)
 
         # avoid broadcast materialization
-        month_emb = (token_emb_seq * day_w).sum(dim=3)
+        aggregated_month_embed = (token_emb_seq * day_w).sum(dim=3)
 
-        month_tokens = month_tokens + month_emb
+        month_tokens = month_tokens + aggregated_month_embed
 
         z = month_tokens.reshape(B * HW, M, C)
 
@@ -617,7 +627,6 @@ def __init__(
         in_chans=1,
         embed_dim=128,
         patch_size=(1, 4, 4),
-        max_months=12,
         hidden=256,
         overlap=1,
         spatial_depth=2,
@@ -633,7 +642,6 @@ def __init__(
             in_chans: Number of input channels (e.g., 1 for SST, additional channels possible)
             embed_dim: Dimension of the patch embedding
             patch_size: Tuple of (T, H, W) patch sizes for temporal and spatial patching
-            max_months: Maximum number of months for temporal positional encoding
             hidden: Hidden dimension used in the decoder
             overlap: Overlap for deconvolution in the decoder
             max_H: Maximum spatial height for 2D positional encoding
@@ -660,7 +668,6 @@ def __init__(
         )
         self.temporal = TemporalAttentionAggregator(
             embed_dim=embed_dim,
-            max_months=max_months,
             dropout=dropout,
         )
         self.geo_embedding = GeoPositionScaleEmbedding(

climanet/utils.py

@@ -120,6 +120,7 @@ def add_month_day_dims(
         .unstack(time_dim)
         .reindex(T=np.arange(1, 32), M=month_keys)
     )
+
     # Force dim order: (M, T, H, W) (and keep any other non-time dims after M,T)
     other_dims = [d for d in daily_ts.dims if d != time_dim]  # e.g. ["H", "W"]
     daily_indexed = daily_indexed.transpose("M", "T", *other_dims)
@@ -148,12 +149,14 @@ def add_month_day_dims(
         .reindex(T=np.arange(1, 32), M=month_keys)
     )
 
-    # determine day-of-year (doy) [and hour-of-day (hod) if applicable], fill NaT with 0 inplace
+    # month-of_year (moy), day-of-year (doy) [and hour-of-day (hod) if applicable], fill NaT with 0 inplace
     # here we choose to use the tropical year length (365.2422 day, which we round to 365.24) as the
     # period to return to the position of the sun relative to the Earth
+    moy_period = 12.0
     doy_period = 365.24
     hod_period = 24.0
 
+    moy = time_indexed.dt.month.fillna(0)
     doy = time_indexed.dt.dayofyear.fillna(0)
 
     if "hour" in dir(time_indexed.dt):
@@ -162,13 +165,14 @@ def add_month_day_dims(
         hod = xr.zeros_like(doy)
 
     # create phase from day and hod
+    moy_phase = 2 * np.pi * (moy - 1.0) / moy_period
     doy_phase = 2 * np.pi * doy / doy_period
     hod_phase = 2 * np.pi * hod / hod_period
 
-    # Stack cyclic encodings into time_features (M,T,2)
-    time_features = xr.concat([doy_phase, hod_phase], dim="feature").transpose(
-        "M", "T", "feature"
-    )
+    # Stack cyclic encodings into time_features (M,T,3)
+    time_features = xr.concat(
+        [moy_phase, doy_phase, hod_phase], dim="feature"
+    ).transpose("M", "T", "feature")
 
     return daily_indexed, monthly_m, padded_days_mask, time_features
 
@@ -327,21 +331,24 @@ def add_month_hour_dims(
         .reindex(T=np.arange(1, 745), M=month_keys)
     )
 
-    # Determine day-of-year (doy) and hour-of-day (hod)
+    # Determine month-of-year, day-of-year (doy) and hour-of-day (hod)
+    moy_period = 12.0
     doy_period = 365.24
     hod_period = 24.0
 
+    moy = time_indexed.dt.month.fillna(0)
     doy = time_indexed.dt.dayofyear.fillna(0)
     hod = time_indexed.dt.hour.fillna(0)
 
-    # Create phase from day and hour
+    # Create phase from month, day and hour
+    moy_phase = 2 * np.pi * (moy - 1.0) / moy_period
     doy_phase = 2 * np.pi * doy / doy_period
     hod_phase = 2 * np.pi * hod / hod_period
 
-    # Stack cyclic encodings into time_features (M, T, 2)
-    time_features = xr.concat([doy_phase, hod_phase], dim="feature").transpose(
-        "M", "T", "feature"
-    )
+    # Stack cyclic encodings into time_features (M, T, 3)
+    time_features = xr.concat(
+        [moy_phase, doy_phase, hod_phase], dim="feature"
+    ).transpose("M", "T", "feature")
 
     return hourly_indexed, monthly_m, padded_hours_mask, time_features