average pool

Author: michielkallenberg

cybench/conf/model/torch_model/temporal_encoder/tokenizer/avg_pool.yaml

@@ -0,0 +1,16 @@
+
+_target_: cybench.models.torch.model_components.temporal_encoder.AvgPoolTokenizer
+
+in_dim: ${model.torch_model.temporal_in_dim}
+embed_dim: ${model.torch_model.embed_dim}
+
+# Non-overlapping per-channel mean pool (dekadal when patch_size=10).
+patch_size: 10
+eos_anchor: true
+
+_search_:
+  patch_size:
+    type: int
+    low: 5
+    high: 15
+    step: 5

cybench/models/torch/model_components/temporal_encoder.py

@@ -2,16 +2,16 @@
 
 import torch
 import torch.nn as nn
-from typing import Optional
+import torch.nn.functional as F
 
 """
 Temporal Encoder Module
 =======================
 
 This module implements a flexible temporal-encoding pipeline composed of:
 1. **Tokenizer**: Downsamples raw temporal features and maps them into a fixed
-   embedding dimension (`embed_dim`). Implemented via a high–receptive-field
-   Conv1d block (`ConvTokenizer`).
+   embedding dimension (`embed_dim`). Options include a learnable Conv1d block
+   (`ConvTokenizer`) or a fixed per-channel mean pool (`AvgPoolTokenizer`).
 2. **Processor**: Applies temporal feature transformation without changing the
    embedding dimensionality. Two processor families are provided:
        - `CNNProcessor`: Lightweight temporal convolutions with preserved shape.
@@ -143,6 +143,65 @@ def forward(self,
             return z
 
 
+class _FixedChannelProjection(nn.Module):
+    """Pad or truncate channels without learnable weights."""
+
+    def __init__(self, in_dim: int, embed_dim: int):
+        super().__init__()
+        self.in_dim = in_dim
+        self.embed_dim = embed_dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.in_dim == self.embed_dim:
+            return x
+        if self.in_dim < self.embed_dim:
+            return F.pad(x, (0, self.embed_dim - self.in_dim))
+        return x[..., : self.embed_dim]
+
+
+class AvgPoolTokenizer(nn.Module):
+    """
+    Fixed temporal tokenizer: non-overlapping per-channel mean over patch_size days.
+
+  Mimics dekadal (or N-day) mean aggregation on the model side. No learnable
+    weights. When ``eos_anchor`` is True (default), the earliest days are dropped
+    so the last pool window ends at the final timestep (EOS-aligned windows).
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        embed_dim: int,
+        patch_size: int = 10,
+        eos_anchor: bool = True,
+    ):
+        super().__init__()
+        if patch_size < 1:
+            raise ValueError(f"patch_size must be >= 1, got {patch_size}")
+        self.patch_size = patch_size
+        self.eos_anchor = eos_anchor
+        self.pool = nn.AvgPool1d(kernel_size=patch_size, stride=patch_size)
+        self.proj = _FixedChannelProjection(in_dim, embed_dim)
+
+    def _trim_to_eos_anchor(self, x: torch.Tensor) -> torch.Tensor:
+        remainder = x.shape[1] % self.patch_size
+        if remainder == 0:
+            return x
+        return x[:, remainder:, :]
+
+    def forward(self, x: torch.Tensor, doys: torch.Tensor) -> torch.Tensor:
+        del doys
+        if self.eos_anchor:
+            x = self._trim_to_eos_anchor(x)
+        if x.shape[1] < self.patch_size:
+            raise ValueError(
+                f"Sequence length {x.shape[1]} is shorter than patch_size "
+                f"{self.patch_size} after EOS anchoring."
+            )
+        z = self.pool(x.transpose(1, 2)).transpose(1, 2)
+        return self.proj(z)
+
+
 # ----------------------------------------------------------------------
 # 2. Processors
 # ----------------------------------------------------------------------

tests/models/test_temporal_encoder.py

@@ -0,0 +1,112 @@
+import pytest
+import torch
+from hydra import compose, initialize
+from hydra.utils import instantiate
+
+from cybench.datasets.data_factory import DataFactory
+from cybench.evaluation.eval import evaluate_predictions
+from cybench.models.torch.model_components.temporal_encoder import (
+    AvgPoolTokenizer,
+    LSTMProcessor,
+    LastStepPooling,
+    TemporalEncoder,
+)
+from cybench.util.config_utils import adjust_model_cfg_to_dataset, remove_keys
+
+
+def test_avg_pool_tokenizer_shape():
+    batch_size, seq_len, in_dim, embed_dim, patch_size = 2, 25, 8, 8, 10
+    tokenizer = AvgPoolTokenizer(
+        in_dim=in_dim, embed_dim=embed_dim, patch_size=patch_size
+    )
+    x = torch.randn(batch_size, seq_len, in_dim)
+    doys = torch.randint(1, 366, (batch_size, seq_len))
+
+    out = tokenizer(x, doys)
+    assert out.shape == (batch_size, 2, embed_dim)
+
+
+def test_avg_pool_tokenizer_channel_means_eos_anchored():
+    patch_size = 10
+    tokenizer = AvgPoolTokenizer(in_dim=1, embed_dim=1, patch_size=patch_size)
+    x = torch.arange(25, dtype=torch.float32).view(1, 25, 1)
+    doys = torch.zeros(1, 25, dtype=torch.int16)
+
+    out = tokenizer(x, doys)
+
+    assert out.shape == (1, 2, 1)
+    assert out[0, 0, 0].item() == pytest.approx(torch.arange(5, 15).float().mean().item())
+    assert out[0, 1, 0].item() == pytest.approx(torch.arange(15, 25).float().mean().item())
+
+
+def test_avg_pool_tokenizer_has_no_trainable_parameters():
+    tokenizer = AvgPoolTokenizer(in_dim=4, embed_dim=8, patch_size=10)
+    assert sum(p.numel() for p in tokenizer.parameters()) == 0
+
+
+def test_avg_pool_tokenizer_pads_channels_to_embed_dim():
+    tokenizer = AvgPoolTokenizer(in_dim=3, embed_dim=5, patch_size=5)
+    x = torch.ones(1, 10, 3)
+    out = tokenizer(x, torch.zeros(1, 10, dtype=torch.int16))
+
+    assert out.shape == (1, 2, 5)
+    assert torch.all(out[..., 3:] == 0)
+
+
+def test_avg_pool_tokenizer_rejects_short_sequence():
+    tokenizer = AvgPoolTokenizer(in_dim=2, embed_dim=2, patch_size=10)
+    x = torch.randn(1, 8, 2)
+    with pytest.raises(ValueError, match="shorter than patch_size"):
+        tokenizer(x, torch.zeros(1, 8, dtype=torch.int16))
+
+
+def test_avg_pool_temporal_encoder_end_to_end():
+    batch_size, seq_len, in_dim, embed_dim = 2, 30, 6, 6
+    encoder = TemporalEncoder(
+        tokenizer=AvgPoolTokenizer(in_dim=in_dim, embed_dim=embed_dim, patch_size=10),
+        processor=LSTMProcessor(embed_dim=embed_dim, num_layers=1, dropout=0.0),
+        pooling=LastStepPooling(),
+        embed_dim=embed_dim,
+    )
+    temporal = torch.randn(batch_size, seq_len, in_dim)
+    doys = torch.randint(1, 366, (batch_size, seq_len))
+
+    out = encoder(temporal, doys)
+    assert out.shape == (batch_size, embed_dim)
+
+
+def test_lstm_lf_with_avg_pool_tokenizer_runs_on_dataset():
+    overrides = [
+        "dataset/crop=wheat",
+        "dataset.country=NL",
+        "dataset.framework=torch",
+        "dataset.use_cache=false",
+        "model=lstm_lf",
+        "experiment.device=cpu",
+        "model.epochs=1",
+        "model/torch_model/temporal_encoder/tokenizer=avg_pool",
+        "model.torch_model.embed_dim=8",
+        "dataset/temporal=no_aggregate",
+        "~dataset.temporal.sources.ndvi",
+        "~dataset.temporal.sources.soil_moisture",
+        "dataset.temporal.sources.meteo.select=[tmin,tmax,tavg,prec,rad,et0,vpd]",
+    ]
+    with initialize(version_base=None, config_path="../../cybench/conf"):
+        cfg = compose(config_name="config", overrides=overrides)
+
+    dataset = DataFactory(cfg.dataset).build()
+    cfg.model = adjust_model_cfg_to_dataset(cfg.model, dataset)
+    model_cfg = remove_keys(cfg.model, "_search_")
+    model = instantiate(model_cfg)
+
+    even_years = {year for year in dataset.years if year % 2 == 0}
+    odd_years = dataset.years - even_years
+    train_dataset, test_dataset = dataset.split_on_years((even_years, odd_years))
+
+    model.fit(train_dataset)
+    test_preds, _ = model.predict(test_dataset)
+    assert test_preds.shape[0] == len(test_dataset)
+
+    targets = test_dataset.targets
+    evaluation_result = evaluate_predictions(targets, test_preds, cfg.evaluation)
+    assert "normalized_rmse" in evaluation_result