diff --git a/climanet/dataset.py b/climanet/dataset.py index f34ddb7..5cf5306 100644 --- a/climanet/dataset.py +++ b/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, diff --git a/climanet/predict.py b/climanet/predict.py index 97cd2f9..0ad5c23 100644 --- a/climanet/predict.py +++ b/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), diff --git a/climanet/st_encoder_decoder.py b/climanet/st_encoder_decoder.py index 58d9317..70e9901 100644 --- a/climanet/st_encoder_decoder.py +++ b/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,15 +244,13 @@ 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. """ @@ -244,8 +258,8 @@ def __init__(self, embed_dim=128, max_months=12, dropout=0.0): 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( diff --git a/climanet/utils.py b/climanet/utils.py index 65b8ac1..a56ae4a 100644 --- a/climanet/utils.py +++ b/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 diff --git a/notebooks/example_daily.ipynb b/notebooks/example_daily.ipynb index 7adec18..cc1def8 100644 --- a/notebooks/example_daily.ipynb +++ b/notebooks/example_daily.ipynb @@ -128,7 +128,7 @@ "outputs": [], "source": [ "patch_size = (1, 4, 4)\n", - "model = SpatioTemporalModel(patch_size=patch_size, overlap=2, num_months=2, embed_dim=64, dropout=0.2, hidden=64)" + "model = SpatioTemporalModel(patch_size=patch_size, overlap=2, embed_dim=64, dropout=0.2, hidden=64)" ] }, { @@ -158,7 +158,7 @@ " input_da=daily_subset[var_name],\n", " monthly_da=monthly_subset_res[var_name],\n", " land_mask=lsm_subset[\"lsm\"],\n", - " patch_size=spatial_patch_size, # based on the patch_size in model\n", + " patch_size=(1, *spatial_patch_size), # based on the patch_size in model\n", " stride=stride,\n", " sh_embed_dim=96,\n", " sh_order_L = 10,\n", @@ -285,7 +285,7 @@ "# inference on test data, verbose is True\n", "batch_size = 10\n", "predictions = predict_monthly_var(\n", - " model=\"runs/best_model.pth\", dataset=test_dataset, batch_size=batch_size, save_predictions=True, device=device, run_dir=run_dir, dataloader_num_workers=dataloader_num_workers,\n", + " model=f\"{run_dir}/best_model.pth\", dataset=test_dataset, batch_size=batch_size, save_predictions=True, device=device, run_dir=run_dir, dataloader_num_workers=dataloader_num_workers,\n", ")" ] }, @@ -315,7 +315,7 @@ " input_da=daily_subset[var_name],\n", " monthly_da=monthly_subset_res[var_name],\n", " land_mask=lsm_subset[\"lsm\"],\n", - " patch_size=spatial_patch_size, # based on the patch_size in model\n", + " patch_size=(1, *spatial_patch_size), # based on the patch_size in model\n", " sh_embed_dim=96,\n", " sh_order_L = 10,\n", ")\n", diff --git a/notebooks/example_hourly.ipynb b/notebooks/example_hourly.ipynb index 596300d..a213ea3 100644 --- a/notebooks/example_hourly.ipynb +++ b/notebooks/example_hourly.ipynb @@ -123,7 +123,7 @@ "set_seed()\n", "\n", "patch_size = (1, 4, 4)\n", - "model = SpatioTemporalModel(patch_size=patch_size, overlap=2, num_months=2, embed_dim=64, dropout=0.2, hidden=64)" + "model = SpatioTemporalModel(patch_size=patch_size, overlap=2, embed_dim=64, dropout=0.2, hidden=64)" ] }, { @@ -152,7 +152,7 @@ " input_da=daily_subset[var_name],\n", " monthly_da=monthly_subset_res[var_name],\n", " land_mask=lsm_subset[\"lsm\"],\n", - " patch_size=spatial_patch_size, # based on the patch_size in model\n", + " patch_size=(1, *spatial_patch_size), # based on the patch_size in model\n", " stride=stride,\n", " sh_embed_dim=96,\n", " sh_order_L = 10,\n", @@ -303,7 +303,7 @@ " input_da=daily_subset[var_name],\n", " monthly_da=monthly_subset_res[var_name],\n", " land_mask=lsm_subset[\"lsm\"],\n", - " patch_size=spatial_patch_size, # based on the patch_size in model\n", + " patch_size=(1, *spatial_patch_size), # based on the patch_size in model\n", " sh_embed_dim=96,\n", " sh_order_L = 10,\n", " is_hourly=True, \n", diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 0ca5fc2..8b1cf69 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -64,7 +64,7 @@ def test_len_and_shapes(): assert sample["daily_patch"].shape == (1, 1, 31, 2, 2) assert sample["monthly_patch"].shape == (1, 2, 2) assert sample["daily_mask_patch"].shape == (1, 1, 31, 2, 2) - assert sample["daily_timef_patch"].shape == (1, 31, 2) + assert sample["daily_timef_patch"].shape == (1, 31, 3) assert sample["daily_patch"].dtype == torch.float32 assert sample["monthly_patch"].dtype == torch.float32 assert sample["daily_mask_patch"].dtype == torch.bool @@ -114,6 +114,6 @@ def test_time_feature_generation(): sample = dataset[0] expected_time_feature = torch.tensor( - [np.float32(2 * np.pi * 6 / 365.24), np.float32(0.0)] + [np.float32(0.), np.float32(2 * np.pi * 6 / 365.24), np.float32(0.0)] ) assert torch.equal(sample["daily_timef_patch"][0, 5, :], expected_time_feature) diff --git a/tests/test_train.py b/tests/test_train.py index c623b81..35cd560 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -14,7 +14,7 @@ def dummy_batch(): "monthly_patch": torch.rand(1, 2, 40, 40), "daily_mask_patch": torch.rand(1, 1, 2, 31, 40, 40) > 0.5, # boolean mask "land_mask_patch": torch.rand(1, 40, 40) > 0.5, # boolean mask - "daily_timef_patch": torch.rand(1, 2, 31, 2), + "daily_timef_patch": torch.rand(1, 2, 31, 3), "padded_days_mask": torch.rand(1, 2, 31) > 0.5, # boolean mask "scale_feature_patch": torch.rand(1, 10), "geo_pos_embedding_patch": torch.rand(1, 96),