Merge pull request #820 from alan-turing-institute/advection-reaction

Add dataset for advection-diffusion example

Author: ContiPaolo

autoemulate/experimental/exploratory/advection_diffusion_fno.ipynb

@@ -0,0 +1,173 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from autoemulate.simulations.advection_diffusion import AdvectionDiffusion\n",
+    "\n",
+    "\n",
+    "rd = AdvectionDiffusion(n=50, T=80, dt=0.25, return_timeseries=True)\n",
+    "data = rd.forward_samples_spatiotemporal(20)\n",
+    "y = data[\"data\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.imshow(y[0, 50, :, :, 0])\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch.utils.data import DataLoader\n",
+    "from autoemulate.experimental.data.spatiotemporal_dataset import AutoEmulateDataset\n",
+    "\n",
+    "dataset = AutoEmulateDataset(data_path=None, data=data, n_steps_input=1, n_steps_output=1)\n",
+    "batch = next(iter(DataLoader(dataset)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch[\"input_fields\"].shape, batch[\"output_fields\"].shape, batch[\"constant_scalars\"].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch[\"input_fields\"].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.imshow(batch[\"input_fields\"][0, 0, :, :, 0].cpu())\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from autoemulate.experimental.emulators.fno import FNOEmulator\n",
+    "\n",
+    "emulator = FNOEmulator(\n",
+    "    n_modes=(1, 16, 16),\n",
+    "    hidden_channels=16,\n",
+    "    in_channels=3,\n",
+    "    out_channels=1,\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Fit the emulator\n",
+    "emulator.fit(DataLoader(dataset), None)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Predictions\n",
+    "y_pred = emulator.predict(DataLoader(dataset), with_grad=False)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Evaluate\n",
+    "# TODO: add to emulator perhaps as .evaluate()?\n",
+    "import torch\n",
+    "from autoemulate.experimental.emulators.fno import prepare_batch\n",
+    "\n",
+    "y_true = torch.cat(\n",
+    "    [\n",
+    "        prepare_batch(\n",
+    "            batch, channels=(0,), with_constants=True, with_time=True\n",
+    "        )[1]\n",
+    "        for batch in DataLoader(dataset)\n",
+    "    ],\n",
+    "    dim=0\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torchmetrics import R2Score\n",
+    "\n",
+    "R2Score()(y_pred.reshape(-1).detach(), y_true.reshape(-1).detach()).item()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}