Merge pull request #794 from alan-turing-institute/788-reaction-diffusion-dataset

Add dataset for reaction diffusion example (#788)

Author: sgreenbury

autoemulate/emulators/base.py

@@ -8,7 +8,6 @@
 from torch import nn, optim
 from torch.distributions import TransformedDistribution
 from torch.optim.lr_scheduler import ExponentialLR, LRScheduler
-from torch.utils.data import DataLoader
 
 from autoemulate.core.device import TorchDeviceMixin
 from autoemulate.core.types import (
@@ -40,9 +39,9 @@ class Emulator(ABC, ValidationMixin, ConversionMixin, TorchDeviceMixin):
     supports_uq: bool = False
 
     @abstractmethod
-    def _fit(self, x: TensorLike | DataLoader, y: TensorLike | DataLoader | None): ...
+    def _fit(self, x: TensorLike, y: TensorLike): ...
 
-    def fit(self, x: TensorLike | DataLoader, y: TensorLike | DataLoader | None):
+    def fit(self, x: TensorLike, y: TensorLike):
         """Fit the emulator to the provided data."""
         if isinstance(x, TensorLike) and isinstance(y, TensorLike):
             self._check(x, y)
@@ -62,11 +61,6 @@ def fit(self, x: TensorLike | DataLoader, y: TensorLike | DataLoader | None):
 
             # Fit emulator
             self._fit(x, y)
-        elif isinstance(x, DataLoader) and y is None:
-            self._fit(x, y)
-        else:
-            msg = "Invalid input types. Expected pair of TensorLike or DataLoader."
-            raise RuntimeError(msg)
         self.is_fitted_ = True
 
     @abstractmethod
@@ -547,7 +541,7 @@ def loss_func(self, y_pred, y_true):
         """Loss function to be used for training the model."""
         return nn.MSELoss()(y_pred, y_true)
 
-    def _fit(self, x: TensorLike, y: TensorLike):  # type: ignore since this is valid subclass of types
+    def _fit(self, x: TensorLike, y: TensorLike):
         """
         Train a PyTorchBackend model.
 
@@ -671,7 +665,7 @@ class SklearnBackend(DeterministicEmulator):
     def _model_specific_check(self, x: NumpyLike, y: NumpyLike):
         _, _ = x, y
 
-    def _fit(self, x: TensorLike, y: TensorLike):  # type: ignore since this is valid subclass of types
+    def _fit(self, x: TensorLike, y: TensorLike):
         if self.normalize_y:
             y, y_mean, y_std = self._normalize(y)
             self.y_mean = y_mean

autoemulate/emulators/ensemble.py

@@ -67,7 +67,7 @@ def get_tune_params() -> TuneParams:
         """Return a dictionary of hyperparameters to tune."""
         return {}
 
-    def _fit(self, x: TensorLike, y: TensorLike) -> None:  # type: ignore since this is valid subclass of types
+    def _fit(self, x: TensorLike, y: TensorLike) -> None:
         for e in self.emulators:
             e.fit(x, y)
         self.is_fitted_ = True
@@ -248,7 +248,7 @@ def get_tune_params() -> TuneParams:
             "n_samples": [10, 20, 50, 100],
         }
 
-    def _fit(self, x: TensorLike, y: TensorLike) -> None:  # type: ignore since this is valid subclass of types
+    def _fit(self, x: TensorLike, y: TensorLike) -> None:
         # Delegate training to the wrapped model
         self.model.fit(x, y)
         self.is_fitted_ = True

autoemulate/emulators/gaussian_process/exact.py

@@ -183,7 +183,7 @@ def forward(self, x: TensorLike):
             MultivariateNormal(mean, covar)
         )
 
-    def _fit(self, x: TensorLike, y: TensorLike):  # type: ignore since this is valid subclass of types
+    def _fit(self, x: TensorLike, y: TensorLike):
         self.train()
         self.likelihood.train()
 

autoemulate/emulators/transformed/base.py

@@ -374,7 +374,7 @@ def _inv_transform_y_distribution(self, y_t: DistributionLike) -> DistributionLi
         """
         return TransformedDistribution(y_t, [ComposeTransform(self.y_transforms).inv])
 
-    def _fit(self, x: TensorLike, y: TensorLike):  # type: ignore since this is valid subclass of types
+    def _fit(self, x: TensorLike, y: TensorLike):
         # Transform x and y
         x_t = self._transform_x(x)
         y_t = self._transform_y_tensor(y)

…rimental/data/spatio_temporal_dataset.py …erimental/data/spatiotemporal_dataset.pyautoemulate/experimental/data/spatio_temporal_dataset.py renamed to autoemulate/experimental/data/spatiotemporal_dataset.py autoemulate/experimental/data/spatio_temporal_dataset.py renamed to autoemulate/experimental/data/spatiotemporal_dataset.py

@@ -9,12 +9,12 @@ class AutoEmulateDataset(Dataset):
 
     def __init__(
         self,
-        data_path: str,
-        n_steps_input: int,
-        n_steps_output: int,
+        data_path: str | None,
+        data: dict | None = None,
+        n_steps_input: int = 1,
+        n_steps_output: int = 1,
         stride: int = 1,
         # TODO: support for passing data from dict
-        # data: dict | None = None,
         input_channel_idxs: tuple[int, ...] | None = None,
         output_channel_idxs: tuple[int, ...] | None = None,
     ):
@@ -45,20 +45,8 @@ def __init__(
         self.input_channel_idxs = input_channel_idxs
         self.output_channel_idxs = output_channel_idxs
 
-        # TODO: support passing as dict
-        # Load data
-        with h5py.File(data_path, "r") as f:
-            assert "data" in f, "HDF5 file must contain 'data' dataset"
-            self.data: TensorLike = torch.Tensor(f["data"][:])  # type: ignore # [N, T, W, H, C]  # noqa: PGH003
-            print(f"Loaded data shape: {self.data.shape}")
-            # TODO: add the constant scalars
-            self.constant_scalars = (
-                torch.Tensor(f["constant_scalars"][:])  # type: ignore  # noqa: PGH003
-                if "constant_scalars" in f
-                else None
-            )  # [N, C]
-            # TODO: add the constant fields
-            # self.constant_fields = torch.Tensor(f['data'][:])  # [N, W, H, C]
+        # Read or parse data
+        self.read_data(data_path) if data_path is not None else self.parse_data(data)
 
         # Destructured here
         (
@@ -107,14 +95,45 @@ def __init__(
         print(f"Each input sample shape: {self.all_input_fields[0].shape}")
         print(f"Each output sample shape: {self.all_output_fields[0].shape}")
 
+    def read_data(self, data_path: str):
+        """Read data.
+
+        By default assumes HDF5 format in `data_path` with correct shape and fields.
+        """
+        # TODO: support passing as dict
+        # Load data
+        self.data_path = data_path
+        with h5py.File(self.data_path, "r") as f:
+            assert "data" in f, "HDF5 file must contain 'data' dataset"
+            self.data: TensorLike = torch.Tensor(f["data"][:])  # type: ignore # [N, T, W, H, C]  # noqa: PGH003
+        print(f"Loaded data shape: {self.data.shape}")
+        # TODO: add the constant scalars
+        self.constant_scalars = (
+            torch.Tensor(f["constant_scalars"][:])  # type: ignore  # noqa: PGH003
+            if "constant_scalars" in f
+            else None
+        )  # [N, C]
+        # TODO: add the constant fields
+        # self.constant_fields = torch.Tensor(f['data'][:])  # [N, W, H, C]
+
+    def parse_data(self, data: dict | None):
+        """Parse data from a dictionary."""
+        if data is not None:
+            self.data = data["data"]
+            self.constant_scalars = data.get("constant_scalars", None)
+            self.constant_fields = data.get("constant_fields", None)
+            return
+        msg = "No data provided to parse."
+        raise ValueError(msg)
+
     def __len__(self):  # noqa: D105
         return len(self.all_input_fields)
 
     def __getitem__(self, idx):  # noqa: D105
         return {
             "input_fields": self.all_input_fields[idx],
             "output_fields": self.all_output_fields[idx],
-            # "constant_scalars": self.all_constant_scalars[idx],
+            "constant_scalars": self.all_constant_scalars[idx],
             # TODO: add this
             # "constant_fields": self.all_constant_fields[idx],
         }

autoemulate/experimental/emulators/fno.py

@@ -1,6 +1,6 @@
 import torch
-from autoemulate.core.types import TensorLike
-from autoemulate.emulators.base import PyTorchBackend
+from autoemulate.core.types import OutputLike, TensorLike
+from autoemulate.experimental.emulators.spatiotemporal import SpatioTemporalEmulator
 from neuralop.models import FNO
 from torch.utils.data import DataLoader
 
@@ -41,14 +41,23 @@ def prepare_batch(sample, channels=(0,), with_constants=True, with_time=False):
     return x, y
 
 
-class FNOEmulator(PyTorchBackend):
+class FNOEmulator(SpatioTemporalEmulator):
     """An FNO emulator."""
 
-    def __init__(self, x, y, *args, **kwargs):
-        _, _ = x, y
+    def __init__(self, x=None, y=None, *args, **kwargs):
+        _, _ = x, y  # Unused
+        # Ensure parent initialisers run before creating nn.Module attributes
+        super().__init__()
         self.model = FNO(**kwargs)
+        self.optimizer = torch.optim.Adam(self.model.parameters())
 
-    def _fit(self, x: DataLoader, y: DataLoader | None):  # type: ignore  # noqa: PGH003
+    @staticmethod
+    def is_multioutput() -> bool:  # noqa: D102
+        return True
+
+    def _fit(self, x: TensorLike | DataLoader, y: TensorLike | None = None):
+        assert isinstance(x, DataLoader), "x currently must be a DataLoader"
+        assert y is None, "y currently must be None"
         channels = (0,)  # Which channel to use
         for idx, batch in enumerate(x):
             # Prepare input with constants
@@ -73,5 +82,16 @@ def forward(self, x: TensorLike):
         """Forward pass."""
         return self.model(x)
 
-    def _predict(self, x, with_grad):
-        return super()._predict(x, with_grad)
+    def _predict(self, x: TensorLike | DataLoader, with_grad: bool) -> OutputLike:
+        assert isinstance(x, DataLoader), "x currently must be a DataLoader"
+        with torch.set_grad_enabled(with_grad):
+            channels = (0,)  # Which channel to use
+            all_preds = []
+            for _, batch in enumerate(x):
+                # Prepare input with constants
+                x, _ = prepare_batch(
+                    batch, channels=channels, with_constants=True, with_time=True
+                )
+                out = self(x)
+                all_preds.append(out)
+            return torch.cat(all_preds)

autoemulate/experimental/emulators/spatiotemporal.py

@@ -0,0 +1,77 @@
+from abc import abstractmethod
+
+from autoemulate.core.types import OutputLike, TensorLike
+from autoemulate.emulators.base import PyTorchBackend
+from torch.utils.data import DataLoader
+
+
+class SpatioTemporalEmulator(PyTorchBackend):
+    """A spatio-temporal backend for emulators."""
+
+    def fit(self, x: TensorLike | DataLoader, y: TensorLike | None = None):
+        """Train a spatio-temporal emulator.
+
+        Parameters
+        ----------
+        x: TensorLike | DataLoader
+            Input features as `TensorLike` or `DataLoader`.
+        y: OutputLike | None
+            Target values (not needed if x is a DataLoader).
+
+        """
+        if isinstance(x, TensorLike) and isinstance(y, TensorLike):
+            return super().fit(x, y)
+        if isinstance(x, DataLoader) and y is None:
+            return self._fit(x, y)
+        msg = "Invalid input types. Expected pair of TensorLike or DataLoader only."
+        raise RuntimeError(msg)
+
+    @abstractmethod
+    def _fit(self, x: TensorLike | DataLoader, y: TensorLike | None = None): ...
+
+    def predict(
+        self, x: TensorLike | DataLoader, with_grad: bool = False
+    ) -> OutputLike:
+        """Predict the output for the given input.
+
+        Parameters
+        ----------
+        x: TensorLike | DataLoader
+            Input `TensorLike` or `DataLoader` to make predictions for.
+        with_grad: bool
+            Whether to enable gradient calculation. Defaults to False.
+
+        Returns
+        -------
+        OutputLike
+            The emulator predicted spatiotemporal output.
+        """
+        if isinstance(x, TensorLike):
+            return super().predict(x, with_grad)
+        return self._predict(x, with_grad)
+
+    @abstractmethod
+    def _predict(self, x: TensorLike | DataLoader, with_grad: bool) -> OutputLike: ...
+
+    # TODO: add method for rollout predictions
+    # def predict_rollout(self, x: DataLoader, timesteps: int = 1) -> OutputLike:
+    #     """
+    #     Predict the output for the given input, rolling out for a number of timesteps.
+
+    #     Parameters
+    #     ----------
+    #     x: DataLoader
+    #         Input `DataLoader` to make predictions for.
+    #     timesteps: int
+    #         Number of timesteps to rollout for. Defaults to 1.
+
+    #     Returns
+    #     -------
+    #     OutputLike
+    #         The emulator predicted spatiotemporal output.
+    #     """
+
+    #     # Start at t=0 x_0
+    #     # model predicts x_1 given x_0
+    #     # then model predicts x_2 given model's predicted x_1
+    #     # then model predicts x_3 given model's predicted x_2