Fix FNO and rollout

- Update FNO and EncoderProcessorDecoder to expect no temporal dim
- Update return from rollout to stack on new first dimension for rollout windows

Author: sgreenbury

notebooks/00_exploration.ipynb

@@ -28,13 +28,11 @@
    "outputs": [],
    "source": [
     "\n",
-    "from autoemulate.simulations.advection_diffusion import AdvectionDiffusion\n",
+    "from autoemulate.simulations.reaction_diffusion import ReactionDiffusion as Sim\n",
     "\n",
-    "sim = AdvectionDiffusion(return_timeseries=True, log_level=\"error\")\n",
+    "sim = Sim(return_timeseries=True, log_level=\"error\")\n",
     "\n",
-    "def generate_split(\n",
-    "    simulator: AdvectionDiffusion, n_train: int = 4, n_valid: int = 2, n_test: int = 2\n",
-    "):\n",
+    "def generate_split(simulator: Sim, n_train: int = 1, n_valid: int = 1, n_test: int = 1):\n",
     "    \"\"\"Generate training, validation, and test splits from the simulator.\"\"\"\n",
     "    train = simulator.forward_samples_spatiotemporal(n_train)\n",
     "    valid = simulator.forward_samples_spatiotemporal(n_valid)\n",
@@ -62,8 +60,10 @@
    "source": [
     "from auto_cast.data.datamodule import SpatioTemporalDataModule\n",
     "\n",
+    "n_steps_input = 4\n",
+    "n_steps_output = 1\n",
     "datamodule = SpatioTemporalDataModule(\n",
-    "    data=combined_data, data_path=None, n_steps_input=4, n_steps_output=1, batch_size=16\n",
+    "    data=combined_data, data_path=None, n_steps_input=n_steps_input, n_steps_output=n_steps_output, batch_size=16\n",
     ")"
    ]
   },
@@ -100,8 +100,13 @@
     "from auto_cast.models.encoder_processor_decoder import EncoderProcessorDecoder\n",
     "from auto_cast.nn.fno import FNOProcessor\n",
     "\n",
+    "batch = next(iter(datamodule.train_dataloader()))\n",
+    "n_channels = batch.input_fields.shape[-1]\n",
     "processor = FNOProcessor(\n",
-    "    in_channels=1, out_channels=1, n_modes=(16, 16, 1), hidden_channels=64\n",
+    "    in_channels=n_channels * n_steps_input,\n",
+    "    out_channels=n_channels * n_steps_output,\n",
+    "    n_modes=(16, 16),\n",
+    "    hidden_channels=64,\n",
     ")\n",
     "encoder = PermuteConcat(with_constants=False)\n",
     "decoder = ChannelsLast()\n",
@@ -113,30 +118,41 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "id": "8",
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "model(batch).shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9",
+   "metadata": {},
    "source": [
     "### Run trainer\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
     "import lightning as L\n",
     "\n",
-    "device = \"mps\"  # \"cpu\"\n",
-    "trainer = L.Trainer(max_epochs=5, accelerator=device, log_every_n_steps=10)\n",
+    "# device = \"mps\"  # \"cpu\"\n",
+    "device = \"cpu\"\n",
+    "trainer = L.Trainer(max_epochs=1, accelerator=device, log_every_n_steps=10)\n",
     "trainer.fit(model, datamodule.train_dataloader(), datamodule.val_dataloader())"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "10",
+   "id": "11",
    "metadata": {},
    "source": [
     "### Run the evaluation"
@@ -145,12 +161,75 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "11",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
     "trainer.test(model, datamodule.test_dataloader())"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13",
+   "metadata": {},
+   "source": [
+    "### Example rollout"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# A single element is the full trajectory\n",
+    "batch = next(iter(datamodule.rollout_test_dataloader()))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# First n_steps_input are inputs\n",
+    "print(batch.input_fields.shape)\n",
+    "# Remaining n_steps_output are outputs\n",
+    "print(batch.output_fields.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "16",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Run rollout on one trajectory\n",
+    "preds, trues = model.rollout(batch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "17",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(preds.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "18",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(trues.shape)\n"
+   ]
   }
  ],
  "metadata": {

src/auto_cast/decoders/channels_last.py

@@ -9,4 +9,4 @@ class ChannelsLast(Decoder):
 
     def forward(self, x: Tensor) -> Tensor:
         """Forward pass through the ChannelsLast decoder."""
-        return rearrange(x, "b c t w h -> b t w h c")
+        return rearrange(x, "b c w h -> b 1 w h c")

src/auto_cast/encoders/permute_concat.py

@@ -26,4 +26,4 @@ def forward(self, batch: Batch) -> Tensor:
             scalars = rearrange(scalars, "b c -> b c 1 1 1")
             scalars = scalars.expand(b, -1, t, w, h)
             x = torch.cat([x, scalars], dim=1)
-        return x
+        return rearrange(x, "b c t w h -> b (c t) w h")

src/auto_cast/processors/rollout.py

@@ -41,9 +41,19 @@ def rollout(self, batch: BatchT) -> RolloutOutput:
 
             current_batch = self._advance_batch(current_batch, next_inputs, self.stride)
 
-        predictions = torch.stack(pred_outs)
+        # Stack along a new axis after batch representing number of rollout windows R
+        # Each window R contains n_steps_output time steps T.
+        # For example with:
+        # - batch size B=16
+        # - rollout windows R=10
+        # - n_steps_output T=2 per window,
+        # - spatial dimensions W=16, H=8
+        # - channels C=2
+        # The output shapes will be:
+        # (B, R, T, W, H, C) = (16, 10, 2, 16, 8, 2)
+        predictions = torch.stack(pred_outs, dim=1)  # (B, R, T, spatial, C)
         if true_outs:
-            return predictions, torch.stack(true_outs)
+            return predictions, torch.stack(true_outs, dim=1)  # (B, R, T, spatial, C)
         return predictions, None
 
     @abstractmethod