Store spatial domain with the model

scripts/backtest.py

@@ -43,6 +43,7 @@
 from sat_pred.dataset import SatelliteDataset, TimePeriod
 from sat_pred.load_model import get_model_from_checkpoints, get_model_from_huggingface
 from sat_pred.predictions import PREDICTION_DIMS, PREDICTION_VAR_NAME, prediction_coords
+from sat_pred.spatial import SpatialGrid
 
 
 
@@ -82,6 +83,7 @@ def __init__(
         self,
         model: torch.nn.Module,
         data_config: dict,
+        spatial_grid: SpatialGrid,
         model_address: str,
         device: torch.device,
     ) -> None:
@@ -90,6 +92,8 @@ def __init__(
         Args:
             model: The trained torch model
             data_config: The config of the data the model was trained on
+            spatial_grid: The grid the model was trained on, for the dataset to select its
+                inputs onto
             model_address: Where the model was loaded from. Saved onto the predictions, so a store
                 records which model made it
             device: The torch device to run the model on
@@ -105,6 +109,7 @@ def __init__(
         self.model = model.to(device).eval()
         self.device = device
         self.model_address = model_address
+        self.spatial_grid = spatial_grid
 
         # The shape of the samples the model was trained on. These are unpacked from the data
         # config rather than the model config so the backtest samples are built exactly the way the
@@ -126,8 +131,10 @@ def from_checkpoint(cls, checkpoint_dir_path: str, device: torch.device) -> "MLM
             checkpoint_dir_path: Path of the checkpoint directory
             device: The torch device to run the model on
         """
-        model, _, data_config, _ = get_model_from_checkpoints(checkpoint_dir_path, val_best=True)
-        return cls(model, data_config, checkpoint_dir_path, device)
+        model, _, data_config, spatial_grid, _ = get_model_from_checkpoints(
+            checkpoint_dir_path, val_best=True
+        )
+        return cls(model, data_config, spatial_grid, checkpoint_dir_path, device)
 
     @classmethod
     def from_huggingface(cls, repo_id: str, revision: str, device: torch.device) -> "MLModel":
@@ -142,8 +149,8 @@ def from_huggingface(cls, repo_id: str, revision: str, device: torch.device) ->
             revision: The commit hash of the model to download
             device: The torch device to run the model on
         """
-        model, data_config = get_model_from_huggingface(repo_id, revision)
-        return cls(model, data_config, f"{repo_id}@{revision}", device)
+        model, data_config, spatial_grid = get_model_from_huggingface(repo_id, revision)
+        return cls(model, data_config, spatial_grid, f"{repo_id}@{revision}", device)
 
     @torch.no_grad()
     def __call__(self, X: torch.Tensor) -> torch.Tensor:
@@ -267,6 +274,7 @@ def __init__(
         history_mins: int,
         sample_freq_mins: int,
         channels,
+        spatial_grid: SpatialGrid,
     ):
         """A torch Dataset for loading model inputs and the init-time they are a forecast from
 
@@ -279,6 +287,7 @@ def __init__(
             channels: The channels the model was trained on, in order, and the constants each one
                 is normalised with. Either a mapping of channel name to constants, or the path of a
                 YAML file holding one
+            spatial_grid: The grid the model was trained on, which the samples are selected onto
         """
 
         # No target is loaded - the backtest only makes predictions, it does not score them
@@ -289,6 +298,7 @@ def __init__(
             forecast_mins=0,
             sample_freq_mins=sample_freq_mins,
             channels=channels,
+            spatial_grid=spatial_grid,
         )
 
         # Only forecast on the half hour
@@ -520,6 +530,7 @@ def backtest(
         history_mins=model.history_mins,
         sample_freq_mins=model.sample_freq_mins,
         channels=model.channels,
+        spatial_grid=model.spatial_grid,
     )
 
     # Without this the run would quietly do nothing and write no store at all, which is easy to

scripts/push_checkpoint_to_huggingface.py

@@ -23,8 +23,10 @@
     MODEL_CARD_NAME,
     MODEL_CONFIG_NAME,
     PYTORCH_WEIGHTS_NAME,
+    SPATIAL_GRID_NAME,
 )
 from sat_pred.load_model import get_model_from_checkpoints
+from sat_pred.spatial import SpatialGrid
 
 from pathlib import Path
 
@@ -46,6 +48,7 @@ def save_model_to_huggingface(
     save_directory: str,
     model_config: dict,
     data_config: dict,
+    spatial_grid: SpatialGrid,
     wandb_repo: str,
     wandb_id: str,
     experiment_config_path: str | None = None,
@@ -65,6 +68,8 @@ def save_model_to_huggingface(
             Model configuration specified as a key/value dictionary.
         data_config:
             Data configuration the model was trained with, specified as a key/value dictionary.
+        spatial_grid:
+            The grid the model was trained on.
         wandb_repo: Identifier of the repo on wandb.
         wandb_id: Identifier of the model on wandb.
         experiment_config_path: Path to the full hydra config of the training run, if it was saved.
@@ -90,6 +95,11 @@ def save_model_to_huggingface(
     with open(save_directory / DATA_CONFIG_NAME, 'w') as outfile:
         yaml.dump(data_config, outfile, default_flow_style=False)
 
+    # Save the grid the model was trained on. Unlike the experiment config below this is not
+    # optional - without it nothing downstream can check that an input covers the area the model
+    # was trained on
+    spatial_grid.save(save_directory / SPATIAL_GRID_NAME)
+
     # Save the full config of the training run, if it was saved with the checkpoint
     if experiment_config_path is not None:
         shutil.copyfile(experiment_config_path, save_directory / FULL_CONFIG_NAME)
@@ -168,9 +178,8 @@ def push_to_huggingface(
         raise ValueError(f"Could not find wandb run '{wandb_id}' within {wandb_repo}")
 
     # Load the model
-    model, model_config, data_config, experiment_config_path = get_model_from_checkpoints(
-        checkpoint_dir_path,
-        val_best=val_best
+    model, model_config, data_config, spatial_grid, experiment_config_path = (
+        get_model_from_checkpoints(checkpoint_dir_path, val_best=val_best)
     )
 
     # Push to hub
@@ -185,6 +194,7 @@ def push_to_huggingface(
         save_directory=model_output_dir,
         model_config=model_config,
         data_config=data_config,
+        spatial_grid=spatial_grid,
         experiment_config_path=experiment_config_path,
         wandb_repo=wandb_repo,
         wandb_id=wandb_id,

src/sat_pred/constants.py

@@ -10,3 +10,4 @@
 FULL_CONFIG_NAME = "full_experiment_config.yaml"
 PYTORCH_WEIGHTS_NAME = "model.safetensors"
 MODEL_CARD_NAME = "README.md"
+SPATIAL_GRID_NAME = "spatial_grid.npz"

src/sat_pred/dataset.py

@@ -21,6 +21,7 @@
 from ocf_data_sampler.select.find_contiguous_time_periods import find_contiguous_t0_periods
 from ocf_data_sampler.torch_datasets.pvnet_dataset import PickleCacheMixin
 from sat_pred.channels import ChannelConfigInput, parse_channel_config
+from sat_pred.spatial import SpatialGrid
 from sat_pred.xr_tensorstore import open_zarr_paths
 
 
@@ -137,6 +138,7 @@ def __init__(
         forecast_mins: int,
         sample_freq_mins: int,
         channels: ChannelConfigInput,
+        spatial_grid: SpatialGrid | None = None,
         preshuffle: bool = False,
         seed: int | None = None,
     ):
@@ -154,6 +156,8 @@ def __init__(
                 the constants each one is clipped and z-scored with. Either a mapping of channel
                 name to constants, or the path of a YAML file holding one - see
                 `configs/datamodule/channels/`
+            spatial_grid: The grid a model was trained on, if the samples should be selected onto
+                it. `None` takes whatever area the store holds
             preshuffle (bool): Whether to shuffle the data - useful for validation.
                 Defaults to False.
             seed (int | None): Seed used to shuffle the data if `preshuffle` is True. Set this to
@@ -172,6 +176,11 @@ def __init__(
         # the store holds them in, so a model always reads the same channel at the same index
         da = da.sel(channel=self.channel_config.names)
 
+        # Selecting the model's own grid out of the store, so an archive covering a different area
+        # is cropped to what the model was trained on, or refused if it cannot cover it
+        if spatial_grid is not None:
+            da = spatial_grid.select(da)
+
         # Convert the satellite data to the given time frequency by selection
         mask = np.mod(da.time_utc.dt.minute, sample_freq_mins) == 0
         da = da.sel(time_utc=mask)
@@ -198,6 +207,11 @@ def __init__(
         self.da = da
         self.t0_times = t0_times
 
+    @property
+    def spatial_grid(self) -> SpatialGrid:
+        """The grid the samples are taken from"""
+        return SpatialGrid.from_dataarray(self.da)
+
     def __setstate__(self, state: dict) -> None:
         """Rebuild a dataset sent to a dataloader worker
 
@@ -343,6 +357,11 @@ def __init__(
         # Parsed once here so a bad config fails before any data is loaded
         self.channel_config = parse_channel_config(channels)
 
+    @property
+    def spatial_grid(self) -> SpatialGrid:
+        """The grid the training samples are taken from"""
+        return self._get_dataset("train").spatial_grid
+
     def _make_dataset(
         self, time_periods: list[TimePeriod], preshuffle: bool = False
     ) -> SatelliteDataset:

src/sat_pred/load_model.py

@@ -10,7 +10,6 @@
 device and call `.eval()` themselves.
 """
 
-import os
 from glob import glob
 
 import hydra
@@ -24,7 +23,9 @@
     FULL_CONFIG_NAME,
     MODEL_CONFIG_NAME,
     PYTORCH_WEIGHTS_NAME,
+    SPATIAL_GRID_NAME,
 )
+from sat_pred.spatial import SpatialGrid
 
 
 def _read_yaml_config(path: str) -> dict:
@@ -60,7 +61,7 @@ def get_checkpoint_path(checkpoint_dir_path: str, val_best: bool = True) -> str:
 def get_model_from_checkpoints(
     checkpoint_dir_path: str,
     val_best: bool = True,
-) -> tuple[torch.nn.Module, dict, dict, str]:
+) -> tuple[torch.nn.Module, dict, dict, SpatialGrid, str]:
     """Load a model from its checkpoint directory
 
     Args:
@@ -73,6 +74,7 @@ def get_model_from_checkpoints(
             model: The trained torch model, with the lightning wrapper discarded
             model_config: The config of the torch model
             data_config: The config of the data the model was trained on
+            spatial_grid: The grid the model was trained on
             experiment_config_path: Path to the full hydra config of the training run. This is
                 None for models trained before these configs were saved
     """
@@ -95,15 +97,17 @@ def get_model_from_checkpoints(
 
     data_config = _read_yaml_config(f"{checkpoint_dir_path}/{DATA_CONFIG_NAME}")
 
+    spatial_grid = SpatialGrid.load(f"{checkpoint_dir_path}/{SPATIAL_GRID_NAME}")
+
     experiment_config_path = f"{checkpoint_dir_path}/{FULL_CONFIG_NAME}"
 
-    return model, model_config, data_config, experiment_config_path
+    return model, model_config, data_config, spatial_grid, experiment_config_path
 
 
 def get_model_from_huggingface(
     repo_id: str,
     revision: str,
-) -> tuple[torch.nn.Module, dict]:
+) -> tuple[torch.nn.Module, dict, SpatialGrid]:
     """Load a model from a huggingface model repo
 
     The repo holds what `scripts/push_checkpoint_to_huggingface.py` pushed. Note that its
@@ -125,6 +129,8 @@ def get_model_from_huggingface(
             data_config: The config of the data the model was trained on. This gives the history,
                 forecast horizon, time resolution and channels the model expects, so the inputs
                 built for it always match what it was trained to read
+            spatial_grid: The grid the model was trained on, for the caller to select its inputs
+                onto
     """
 
     download_dir = snapshot_download(repo_id=repo_id, revision=revision)
@@ -141,4 +147,6 @@ def get_model_from_huggingface(
 
     data_config = _read_yaml_config(f"{download_dir}/{DATA_CONFIG_NAME}")
 
-    return model, data_config
+    spatial_grid = SpatialGrid.load(f"{download_dir}/{SPATIAL_GRID_NAME}")
+
+    return model, data_config, spatial_grid

src/sat_pred/spatial.py

@@ -0,0 +1,84 @@
+"""The spatial grid a model was trained on, and selecting model inputs onto it
+
+The training data is a crop of the full satellite disc - for the UK models, 614 x 372 pixels at
+about 3km. Nothing about that crop is recorded in the model config, and the models are fully
+convolutional, so one handed a different crop at inference time runs happily and predicts plausible
+values over the wrong geography. The grid is therefore recorded alongside the model when training
+starts, and the inputs are selected onto it before they reach the model: an archive covering more
+than the model needs is cropped to it, and one which cannot cover it is refused.
+
+The coordinates are recorded as full arrays rather than as an extent and a size. The grid is
+regular, so an extent would determine it, but rebuilding it with `linspace` would not reproduce the
+original floats bit for bit, and the selection is an exact lookup.
+"""
+
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+from numpy.typing import NDArray
+
+
+class SpatialGrid:
+    """The spatial grid a model was trained on"""
+
+    def __init__(
+        self,
+        x_geostationary: NDArray[np.float64],
+        y_geostationary: NDArray[np.float64],
+    ) -> None:
+        """The spatial grid a model was trained on
+
+        Both coordinates must be as `open_sat_data` leaves them. It runs
+        `make_spatial_coords_increasing`, which reverses `x_geostationary` - the stores hold it
+        descending - so a grid recorded on one side of that transform would never match one
+        recorded on the other.
+
+        Args:
+            x_geostationary: The x coordinate of every pixel column
+            y_geostationary: The y coordinate of every pixel row
+        """
+        self.x_geostationary = x_geostationary
+        self.y_geostationary = y_geostationary
+
+    @classmethod
+    def from_dataarray(cls, da: xr.DataArray) -> "SpatialGrid":
+        """The grid some satellite data is on
+
+        Args:
+            da: Satellite data as `open_sat_data` returns it
+        """
+        return cls(da.x_geostationary.values, da.y_geostationary.values)
+
+    @classmethod
+    def load(cls, path: str | Path) -> "SpatialGrid":
+        """Load a grid saved by `save`
+
+        Args:
+            path: Path of the file to load
+        """
+        with np.load(path) as file:
+            return cls(file["x_geostationary"], file["y_geostationary"])
+
+    def save(self, path: str | Path) -> None:
+        """Save the grid so that `load` can read it back
+
+        Numpy's own container rather than one of the YAML configs saved beside it, because two
+        float64 coordinate arrays round-trip through it exactly, without the check having to depend
+        on how a text format happens to render a float.
+
+        Args:
+            path: Path of the file to write. Numpy appends `.npz` if it is not already there
+        """
+        np.savez(path, x_geostationary=self.x_geostationary, y_geostationary=self.y_geostationary)
+
+    def select(self, da: xr.DataArray) -> xr.DataArray:
+        """Select the pixels of this grid out of some satellite data
+
+        Data covering more than the grid is cropped to it. Data which does not cover all of it
+        raises, because a model cannot be run over an area it was never trained on.
+
+        Args:
+            da: Satellite data as `open_sat_data` returns it
+        """
+        return da.sel(x_geostationary=self.x_geostationary, y_geostationary=self.y_geostationary)