Update notebook

Author: sgreenbury

autoemulate/experimental/exploratory/the_well/the_well_models_and_metrics.ipynb

@@ -14,7 +14,7 @@
     "from the_well.data import WellDataModule\n",
     "import matplotlib.pyplot as plt\n",
     "import logging\n",
-    "from autoemulate.experimental.emulators.the_well import TheWellEmulator, TheWellFNO\n",
+    "from autoemulate.experimental.emulators.the_well import TheWellFNO\n",
     "\n",
     "logger = logging.getLogger()\n",
     "\n",
@@ -66,7 +66,7 @@
    "outputs": [],
    "source": [
     "# Fit the model\n",
-    "em._fit(None, None)"
+    "# em._fit(None, None)"
    ]
   },
   {
@@ -89,7 +89,7 @@
    "source": [
     "# Example of predictions with the rollout model from a single batch\n",
     "batch = next(iter(datamodule.val_dataloader()))\n",
-    "em.trainer.rollout_model(em.trainer.model, batch, em.trainer.formatter, train=False)[0].shape"
+    "batch_pred = em.trainer.rollout_model(em.trainer.model, batch, em.trainer.formatter, train=False)[0]"
    ]
   },
   {
@@ -99,7 +99,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "batch = next(iter(datamodule.val_dataloader()))"
+    "batch_pred.shape"
    ]
   },
   {
@@ -109,7 +109,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "batch[\"input_fields\"].shape, batch[\"output_fields\"].shape, "
+    "batch = next(iter(datamodule.val_dataloader()))\n",
+    "batch[\"input_fields\"].shape, batch[\"output_fields\"].shape\n",
+    "# (torch.Size([4, 1, 128, 384, 4]), torch.Size([4, 1, 128, 384, 4]))"
    ]
   },
   {
@@ -119,7 +121,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "batch_rollout = next(iter(datamodule.rollout_val_dataloader()))"
+    "batch_rollout = next(iter(datamodule.rollout_val_dataloader()))\n",
+    "batch_rollout_pred = em.trainer.rollout_model(em.trainer.model, batch_rollout, em.trainer.formatter, train=False)[0]"
    ]
   },
   {
@@ -129,7 +132,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "batch_rollout[\"input_fields\"].shape, batch_rollout[\"output_fields\"].shape, \n"
+    "# Output is the whole time series from the \n",
+    "print(batch_rollout[\"input_fields\"].shape, batch_rollout[\"output_fields\"].shape)\n",
+    "# (torch.Size([1, 1, 128, 384, 4]), torch.Size([1, 100, 128, 384, 4]))\n",
+    "print(batch_rollout_pred.shape)\n",
+    "# torch.Size([1, 10, 128, 384, 4]) : max_rollout_steps (10) from current input"
    ]
   },
   {
@@ -139,7 +146,31 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "em.trainer.max_rollout_steps"
+    "# Iteration over the batches covers the whole validation set since it is based from the\n",
+    "# initial conditions in the validation set\n",
+    "for i, batch in enumerate(datamodule.rollout_val_dataloader()):\n",
+    "    print(f\"Batch {i}\")\n",
+    "    # Output is the whole time series from the \n",
+    "    print(batch_rollout[\"input_fields\"].shape, batch_rollout[\"output_fields\"].shape)\n",
+    "    # (torch.Size([1, 1, 128, 384, 4]), torch.Size([1, 100, 128, 384, 4]))\n",
+    "    batch_rollout_pred = em.trainer.rollout_model(em.trainer.model, batch_rollout, em.trainer.formatter, train=False)[0]\n",
+    "    print(batch_rollout_pred.shape)\n",
+    "    # torch.Size([1, 10, 128, 384, 4]) : max_rollout_steps (10) from current input    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Validation loop\n",
+    "em.trainer.validation_loop(\n",
+    "    datamodule.rollout_val_dataloader(),\n",
+    "    valid_or_test=\"rollout_valid\",\n",
+    "    full=True\n",
+    ")"
    ]
   }
  ],