From 0596e1a4776b7e473240ee74f8f4493642238708 Mon Sep 17 00:00:00 2001 From: James Fulton Date: Wed, 19 Aug 2026 15:36:54 +0000 Subject: [PATCH 1/7] Make class to extract, save, load and apply spatial bounds --- src/sat_pred/constants.py | 1 + src/sat_pred/spatial.py | 84 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) create mode 100644 src/sat_pred/spatial.py diff --git a/src/sat_pred/constants.py b/src/sat_pred/constants.py index ad2d2c4..14f3ecd 100644 --- a/src/sat_pred/constants.py +++ b/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" diff --git a/src/sat_pred/spatial.py b/src/sat_pred/spatial.py new file mode 100644 index 0000000..86693f2 --- /dev/null +++ b/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) From 343dac76dbd1b3d394821dbcecc0ad3233e545e9 Mon Sep 17 00:00:00 2001 From: James Fulton Date: Wed, 19 Aug 2026 15:53:08 +0000 Subject: [PATCH 2/7] Save spatial coords at train start --- src/sat_pred/dataset.py | 19 ++++++++++++ src/sat_pred/train.py | 69 ++++++++++++++++++++++------------------- 2 files changed, 56 insertions(+), 32 deletions(-) diff --git a/src/sat_pred/dataset.py b/src/sat_pred/dataset.py index 155a2e3..39ba188 100644 --- a/src/sat_pred/dataset.py +++ b/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: diff --git a/src/sat_pred/train.py b/src/sat_pred/train.py index ea16b55..8525eb0 100755 --- a/src/sat_pred/train.py +++ b/src/sat_pred/train.py @@ -2,6 +2,7 @@ import os import hydra +from pathlib import Path from lightning.pytorch import ( Callback, LightningDataModule, @@ -18,7 +19,12 @@ import rich.tree from lightning.pytorch.utilities import rank_zero_only -from sat_pred.constants import DATA_CONFIG_NAME, FULL_CONFIG_NAME, MODEL_CONFIG_NAME +from sat_pred.constants import ( + DATA_CONFIG_NAME, + FULL_CONFIG_NAME, + MODEL_CONFIG_NAME, + SPATIAL_GRID_NAME, +) from sat_pred.load_model import get_checkpoint_path, get_model_from_checkpoints from sat_pred.loss import LossFunction @@ -135,16 +141,14 @@ def train(config: DictConfig): if "_target_" in cb_conf: callbacks.append(hydra.utils.instantiate(cb_conf)) + # Instantiate the datamodule + datamodule: LightningDataModule = hydra.utils.instantiate(config.datamodule, _convert_='all') + # Align the wandb id with the checkpoint path # - only works if wandb logger and model checkpoint used - use_wandb_logger = False - for logger in loggers: - if isinstance(logger, WandbLogger): - use_wandb_logger = True - wandb_logger = logger - break - - if use_wandb_logger: + wandb_logger = next((lg for lg in loggers if isinstance(lg, WandbLogger)), None) + + if wandb_logger is not None: # Calling the .experiment property initialises the logger wandb_run = wandb_logger.experiment @@ -156,34 +160,35 @@ def train(config: DictConfig): wandb_run.define_metric("trainer/global_step") wandb_run.define_metric("*", step_metric="trainer/global_step") - for callback in callbacks: - if isinstance(callback, ModelCheckpoint): - # skip for non-rank-0 processes: - # see https://github.com/Lightning-AI/pytorch-lightning/issues/13166#issuecomment-1139765549 - if wandb_logger.version is None: - break + checkpoint_callback = next( + (cb for cb in callbacks if isinstance(cb, ModelCheckpoint)), None + ) - callback.dirpath = "/".join( - callback.dirpath.split("/")[:-1] + [wandb_logger.version] - ) - # Also save model config to this path - os.makedirs(callback.dirpath, exist_ok=True) - OmegaConf.save(config.model, f"{callback.dirpath}/{MODEL_CONFIG_NAME}") + # A version of None means a non-rank-0 process, which must not write any of this: + # see https://github.com/Lightning-AI/pytorch-lightning/issues/13166#issuecomment-1139765549 + if checkpoint_callback is not None and wandb_logger.version is not None: - # Similarly save the data config - OmegaConf.save(config.datamodule, f"{callback.dirpath}/{DATA_CONFIG_NAME}") + dirpath = str(Path(checkpoint_callback.dirpath).with_name(wandb_logger.version)) - # Save the full resolved hydra config to this path and upload it to wandb - full_config_path = f"{callback.dirpath}/{FULL_CONFIG_NAME}" - OmegaConf.save(config, full_config_path, resolve=True) - wandb_logger.experiment.save(full_config_path, base_path=callback.dirpath) + checkpoint_callback.dirpath = dirpath - break + # Also save model config to this path + os.makedirs(dirpath, exist_ok=True) + OmegaConf.save(config.model, f"{dirpath}/{MODEL_CONFIG_NAME}") - # Instantiate the datamodule - datamodule: LightningDataModule = hydra.utils.instantiate(config.datamodule, _convert_='all') - - datamodule.zarr_path = list(datamodule.zarr_path) + # Similarly save the data config + OmegaConf.save(config.datamodule, f"{dirpath}/{DATA_CONFIG_NAME}") + + # And the grid the model is about to be trained on, which the data config does not + # record. Saved now rather than when the model is pushed, because the zarr paths could + # have gone stale by then. This builds the train dataset, so the valid-t0 search + # happens here instead of in `fit` - it is cached, so `fit` reuses it + datamodule.spatial_grid.save(f"{dirpath}/{SPATIAL_GRID_NAME}") + + # Save the full resolved hydra config to this path and upload it to wandb + full_config_path = f"{dirpath}/{FULL_CONFIG_NAME}" + OmegaConf.save(config, full_config_path, resolve=True) + wandb_logger.experiment.save(full_config_path, base_path=dirpath) trainer: Trainer = hydra.utils.instantiate( config.trainer, From 499620240dfc9952d1c8ddb765015ee355bb85a9 Mon Sep 17 00:00:00 2001 From: James Fulton Date: Wed, 19 Aug 2026 16:53:17 +0000 Subject: [PATCH 3/7] Upload spatial coords to hf --- scripts/push_checkpoint_to_huggingface.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/scripts/push_checkpoint_to_huggingface.py b/scripts/push_checkpoint_to_huggingface.py index e995bfe..7c48206 100755 --- a/scripts/push_checkpoint_to_huggingface.py +++ b/scripts/push_checkpoint_to_huggingface.py @@ -23,6 +23,7 @@ MODEL_CARD_NAME, MODEL_CONFIG_NAME, PYTORCH_WEIGHTS_NAME, + SPATIAL_GRID_NAME, ) from sat_pred.load_model import get_model_from_checkpoints @@ -46,6 +47,7 @@ def save_model_to_huggingface( save_directory: str, model_config: dict, data_config: dict, + spatial_grid_path: str, wandb_repo: str, wandb_id: str, experiment_config_path: str | None = None, @@ -65,6 +67,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_path: + Path to 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 +94,12 @@ 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. A checkpoint from before the grid was recorded raises here, and needs the + # grid backfilling into it before it can be pushed + shutil.copyfile(spatial_grid_path, 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) @@ -185,6 +195,7 @@ def push_to_huggingface( save_directory=model_output_dir, model_config=model_config, data_config=data_config, + spatial_grid_path=f"{checkpoint_dir_path}/{SPATIAL_GRID_NAME}", experiment_config_path=experiment_config_path, wandb_repo=wandb_repo, wandb_id=wandb_id, From 0cda7d9d35c2a3fab630ec8229f2cec44d062e54 Mon Sep 17 00:00:00 2001 From: James Fulton Date: Wed, 19 Aug 2026 16:53:27 +0000 Subject: [PATCH 4/7] Clean up --- src/sat_pred/train.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/sat_pred/train.py b/src/sat_pred/train.py index 8525eb0..4e701b4 100755 --- a/src/sat_pred/train.py +++ b/src/sat_pred/train.py @@ -68,10 +68,11 @@ def print_config( """ style = "dim" - tree = rich.tree.Tree("CONFIG", style=style, guide_style=style) + style_kwargs = {"style": style, "guide_style": style} + tree = rich.tree.Tree("CONFIG", **style_kwargs) for field in fields: - branch = tree.add(field, style=style, guide_style=style) + branch = tree.add(field, **style_kwargs) config_section = config.get(field) branch_content = str(config_section) @@ -83,6 +84,19 @@ def print_config( rich.print(tree) +def get_next_instance(sequence, class_type) -> object | None: + """Get the next instance of a class in a sequence, or None if there is none. + + Args: + sequence: A sequence of objects to search through. + class_type: The class type to look for. + """ + for item in sequence: + if isinstance(item, class_type): + return item + return None + + @hydra.main(config_path="../../configs/", config_name="config.yaml", version_base="1.2") def train(config: DictConfig): """Train the model using parameters in the supplied config files. @@ -146,7 +160,7 @@ def train(config: DictConfig): # Align the wandb id with the checkpoint path # - only works if wandb logger and model checkpoint used - wandb_logger = next((lg for lg in loggers if isinstance(lg, WandbLogger)), None) + wandb_logger = get_next_instance(loggers, WandbLogger) if wandb_logger is not None: # Calling the .experiment property initialises the logger @@ -160,20 +174,19 @@ def train(config: DictConfig): wandb_run.define_metric("trainer/global_step") wandb_run.define_metric("*", step_metric="trainer/global_step") - checkpoint_callback = next( - (cb for cb in callbacks if isinstance(cb, ModelCheckpoint)), None - ) + checkpoint_callback = get_next_instance(callbacks, ModelCheckpoint) # A version of None means a non-rank-0 process, which must not write any of this: # see https://github.com/Lightning-AI/pytorch-lightning/issues/13166#issuecomment-1139765549 if checkpoint_callback is not None and wandb_logger.version is not None: dirpath = str(Path(checkpoint_callback.dirpath).with_name(wandb_logger.version)) + os.makedirs(dirpath, exist_ok=True) + # Set the checkpoint callback so it writes to this path checkpoint_callback.dirpath = dirpath # Also save model config to this path - os.makedirs(dirpath, exist_ok=True) OmegaConf.save(config.model, f"{dirpath}/{MODEL_CONFIG_NAME}") # Similarly save the data config From d34a6f2c7c39abd379d25793e360b37fbec4b1a0 Mon Sep 17 00:00:00 2001 From: James Fulton Date: Wed, 19 Aug 2026 17:02:12 +0000 Subject: [PATCH 5/7] Load spatial coords --- src/sat_pred/load_model.py | 11 ++++++++--- tests/conftest.py | 9 +++++++++ tests/test_load_model.py | 9 ++++++--- 3 files changed, 23 insertions(+), 6 deletions(-) diff --git a/src/sat_pred/load_model.py b/src/sat_pred/load_model.py index 44b29fc..50b3bf5 100755 --- a/src/sat_pred/load_model.py +++ b/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: @@ -103,7 +104,7 @@ def get_model_from_checkpoints( 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 +126,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 +144,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 diff --git a/tests/conftest.py b/tests/conftest.py index 19a780d..a37de92 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -22,7 +22,9 @@ FULL_CONFIG_NAME, MODEL_CONFIG_NAME, PYTORCH_WEIGHTS_NAME, + SPATIAL_GRID_NAME, ) +from sat_pred.spatial import SpatialGrid # The channel config the training configs use CHANNELS_PATH = Path(__file__).parents[1] / "configs/datamodule/channels/seviri_rss.yaml" @@ -246,4 +248,11 @@ def huggingface_dir(tmp_path, tiny_model_config) -> str: save_model_as_safetensor(model, str(tmp_path / PYTORCH_WEIGHTS_NAME)) + # The grid of the synthetic stores, with `x_geostationary` ascending as `open_sat_data` leaves + # it rather than descending as the stores hold it + SpatialGrid( + x_geostationary=np.sort(np.arange(IMAGE_SIZE, dtype=np.float64) * -3000), + y_geostationary=np.arange(IMAGE_SIZE, dtype=np.float64) * 3000, + ).save(tmp_path / SPATIAL_GRID_NAME) + return str(tmp_path) diff --git a/tests/test_load_model.py b/tests/test_load_model.py index 00de9be..a864311 100644 --- a/tests/test_load_model.py +++ b/tests/test_load_model.py @@ -11,7 +11,7 @@ from sat_pred.models.simvp_model import SimVP from sat_pred.training_module import TrainingModule -from .conftest import HISTORY_MINS +from .conftest import HISTORY_MINS, IMAGE_SIZE # Stand-ins for the repo and pinned commit a caller would name HF_REPO_ID = "some-org/some-model" @@ -76,8 +76,8 @@ def fake_snapshot_download(repo_id: str, revision: str) -> str: def test_get_model_from_huggingface(downloaded_huggingface_dir): - """A model pushed to huggingface round-trips back to a model and its data config""" - model, data_config = get_model_from_huggingface(HF_REPO_ID, revision=HF_REVISION) + """A model pushed to huggingface round-trips back to a model, data config and grid""" + model, data_config, spatial_grid = get_model_from_huggingface(HF_REPO_ID, HF_REVISION) # The huggingface repo holds the bare model, whereas a checkpoint directory holds the training # module wrapping it - so a loader which conflated the two would fail here @@ -86,6 +86,9 @@ def test_get_model_from_huggingface(downloaded_huggingface_dir): assert data_config["history_mins"] == HISTORY_MINS + assert len(spatial_grid.x_geostationary) == IMAGE_SIZE + assert len(spatial_grid.y_geostationary) == IMAGE_SIZE + # `strict=True` already catches a key mismatch. This catches the weights never being loaded at # all, which leaves the model randomly initialised and forecasting silently wrong weights_path = f"{downloaded_huggingface_dir}/{PYTORCH_WEIGHTS_NAME}" From 2c4910e1228a5538def7aba43f205130bf6f0635 Mon Sep 17 00:00:00 2001 From: James Fulton Date: Wed, 19 Aug 2026 17:31:21 +0000 Subject: [PATCH 6/7] Load spatial coords in backtest --- scripts/backtest.py | 19 +++++++++++++++---- scripts/push_checkpoint_to_huggingface.py | 19 +++++++++---------- src/sat_pred/load_model.py | 7 +++++-- src/sat_pred/train.py | 2 +- tests/conftest.py | 16 ++++++++++------ tests/test_load_model.py | 5 +++-- 6 files changed, 43 insertions(+), 25 deletions(-) diff --git a/scripts/backtest.py b/scripts/backtest.py index 952d1e2..4f46f3d 100755 --- a/scripts/backtest.py +++ b/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 diff --git a/scripts/push_checkpoint_to_huggingface.py b/scripts/push_checkpoint_to_huggingface.py index 7c48206..2a4189c 100755 --- a/scripts/push_checkpoint_to_huggingface.py +++ b/scripts/push_checkpoint_to_huggingface.py @@ -26,6 +26,7 @@ SPATIAL_GRID_NAME, ) from sat_pred.load_model import get_model_from_checkpoints +from sat_pred.spatial import SpatialGrid from pathlib import Path @@ -47,7 +48,7 @@ def save_model_to_huggingface( save_directory: str, model_config: dict, data_config: dict, - spatial_grid_path: str, + spatial_grid: SpatialGrid, wandb_repo: str, wandb_id: str, experiment_config_path: str | None = None, @@ -67,8 +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_path: - Path to the grid the model was trained on. + 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. @@ -96,9 +97,8 @@ def save_model_to_huggingface( # 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. A checkpoint from before the grid was recorded raises here, and needs the - # grid backfilling into it before it can be pushed - shutil.copyfile(spatial_grid_path, save_directory / SPATIAL_GRID_NAME) + # 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: @@ -178,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 @@ -195,7 +194,7 @@ def push_to_huggingface( save_directory=model_output_dir, model_config=model_config, data_config=data_config, - spatial_grid_path=f"{checkpoint_dir_path}/{SPATIAL_GRID_NAME}", + spatial_grid=spatial_grid, experiment_config_path=experiment_config_path, wandb_repo=wandb_repo, wandb_id=wandb_id, diff --git a/src/sat_pred/load_model.py b/src/sat_pred/load_model.py index 50b3bf5..aaf7004 100755 --- a/src/sat_pred/load_model.py +++ b/src/sat_pred/load_model.py @@ -61,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: @@ -74,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 """ @@ -96,9 +97,11 @@ 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( diff --git a/src/sat_pred/train.py b/src/sat_pred/train.py index 4e701b4..25f6ad2 100755 --- a/src/sat_pred/train.py +++ b/src/sat_pred/train.py @@ -118,7 +118,7 @@ def train(config: DictConfig): val_best = config.model.model.val_best # Load the model from the checkpoint - torch_model, model_config, _, _ = get_model_from_checkpoints( + torch_model, model_config, *_ = get_model_from_checkpoints( checkpoint_dir, val_best=val_best ) diff --git a/tests/conftest.py b/tests/conftest.py index a37de92..6bd8383 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -53,6 +53,13 @@ "IR_134", "VIS006", "VIS008", "WV_062", "WV_073", ] +# The grid of the synthetic stores, with `x_geostationary` ascending as `open_sat_data` leaves it +# rather than descending as the stores hold it +SPATIAL_GRID = SpatialGrid( + x_geostationary=np.sort(np.arange(IMAGE_SIZE, dtype=np.float64) * -3000), + y_geostationary=np.arange(IMAGE_SIZE, dtype=np.float64) * 3000, +) + AREA_STRING = json.dumps( { "msg_seviri_rss_3km": { @@ -226,6 +233,8 @@ def checkpoint_dir(tmp_path, training_module_config) -> str: yaml.dump({"seed": 12345, "model": training_module_config}) ) + SPATIAL_GRID.save(tmp_path / SPATIAL_GRID_NAME) + torch.save({"state_dict": training_module.state_dict()}, tmp_path / "epoch=1-step=10.ckpt") return str(tmp_path) @@ -248,11 +257,6 @@ def huggingface_dir(tmp_path, tiny_model_config) -> str: save_model_as_safetensor(model, str(tmp_path / PYTORCH_WEIGHTS_NAME)) - # The grid of the synthetic stores, with `x_geostationary` ascending as `open_sat_data` leaves - # it rather than descending as the stores hold it - SpatialGrid( - x_geostationary=np.sort(np.arange(IMAGE_SIZE, dtype=np.float64) * -3000), - y_geostationary=np.arange(IMAGE_SIZE, dtype=np.float64) * 3000, - ).save(tmp_path / SPATIAL_GRID_NAME) + SPATIAL_GRID.save(tmp_path / SPATIAL_GRID_NAME) return str(tmp_path) diff --git a/tests/test_load_model.py b/tests/test_load_model.py index a864311..c11328b 100644 --- a/tests/test_load_model.py +++ b/tests/test_load_model.py @@ -20,8 +20,8 @@ def test_get_model_from_checkpoints(checkpoint_dir, tiny_model_config): """A saved checkpoint round-trips back to a model and its configs""" - model, model_config, data_config, experiment_config_path = get_model_from_checkpoints( - checkpoint_dir + model, model_config, data_config, spatial_grid, experiment_config_path = ( + get_model_from_checkpoints(checkpoint_dir) ) # The lightning wrapper is discarded, leaving the torch model @@ -33,6 +33,7 @@ def test_get_model_from_checkpoints(checkpoint_dir, tiny_model_config): assert model_config == tiny_model_config assert data_config["history_mins"] > 0 + assert len(spatial_grid.x_geostationary) == IMAGE_SIZE assert experiment_config_path == f"{checkpoint_dir}/{FULL_CONFIG_NAME}" From c2cf7d15a1c8e4d467da6dec65c195b5724a8033 Mon Sep 17 00:00:00 2001 From: James Fulton Date: Wed, 19 Aug 2026 17:39:10 +0000 Subject: [PATCH 7/7] Add tests --- tests/test_dataset.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 8634efd..401e315 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -7,6 +7,7 @@ import pytest from sat_pred.dataset import SatelliteDataModule, SatelliteDataset, find_valid_t0_times +from sat_pred.spatial import SpatialGrid from tests.conftest import ( CHANNELS_PATH, DATA_FREQ_MINS, @@ -138,6 +139,28 @@ def test_a_channel_the_data_does_not_have_is_an_error(sat_zarr_path, channel_con ) +def test_a_dataset_can_be_selected_onto_a_grid(sat_zarr_path): + """Samples are cropped to the grid given, and the dataset reports the grid it is on + + Recording the grid at training and selecting onto it at inference only line up if these two + are inverses of each other. + """ + full_dataset = make_dataset(sat_zarr_path, [PERIOD_1]) + + grid = SpatialGrid( + x_geostationary=full_dataset.da.x_geostationary.values[2:10], + y_geostationary=full_dataset.da.y_geostationary.values[3:12], + ) + + dataset = make_dataset(sat_zarr_path, [PERIOD_1], spatial_grid=grid) + + np.testing.assert_array_equal(dataset.spatial_grid.x_geostationary, grid.x_geostationary) + np.testing.assert_array_equal(dataset.spatial_grid.y_geostationary, grid.y_geostationary) + + X, _ = dataset[0] + assert X.shape == (NUM_CHANNELS, NUM_HISTORY_STEPS, 9, 8) + + def test_samples_are_normalised(sat_zarr_path, channel_config): """Samples come out as per-channel z-scores of the clipped physical values""" dataset = make_dataset(sat_zarr_path, [PERIOD_1])