From 3fbeea0e3d05e1a8c08b25470ca289bbcdbacbf2 Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Wed, 29 Jul 2026 18:46:59 +0000
Subject: [PATCH 1/9] Split generation into input and target

---
 src/ocf_data_sampler/config/__init__.py       |   2 +-
 src/ocf_data_sampler/config/load.py           |   6 +-
 src/ocf_data_sampler/config/model.py          | 159 ++++++++---
 src/ocf_data_sampler/config/save.py           |   6 +-
 .../datasets/pvnet/dataset.py                 | 260 +++++++++---------
 .../datasets/pvnet/loading.py                 |  42 ++-
 .../datasets/pvnet/preprocess.py              | 181 +++++++++---
 src/ocf_data_sampler/datasets/pvnet/sample.py |  36 +--
 .../datasets/pvnet/slicing.py                 |  70 ++---
 .../datasets/pvnet/valid_t0s.py               |  40 +--
 src/ocf_data_sampler/load/__init__.py         |   1 +
 src/ocf_data_sampler/load/generation.py       |   5 -
 src/ocf_data_sampler/load/locations.py        |  47 ++++
 src/ocf_data_sampler/select/dropout.py        |  11 +-
 tests/config/test_config.py                   |  83 ++++--
 tests/config/test_load.py                     |   4 +-
 tests/conftest.py                             |  47 +++-
 tests/datasets/pvnet/test_dataset.py          |  91 ++++--
 tests/datasets/pvnet/test_preprocess.py       |  86 ++++--
 tests/datasets/pvnet/test_sample.py           |  39 +--
 tests/fixtures/configs/pvnet_test_config.yaml |  99 +++----
 tests/load/test_load_generation.py            |   1 -
 tests/load/test_load_locations.py             |  56 ++++
 tests/select/test_dropout.py                  |  23 +-
 24 files changed, 930 insertions(+), 465 deletions(-)
 create mode 100644 src/ocf_data_sampler/load/locations.py
 create mode 100644 tests/load/test_load_locations.py

diff --git a/src/ocf_data_sampler/config/__init__.py b/src/ocf_data_sampler/config/__init__.py
index 78ead638..b4040a48 100644
--- a/src/ocf_data_sampler/config/__init__.py
+++ b/src/ocf_data_sampler/config/__init__.py
@@ -1,5 +1,5 @@
 """Configuration model."""
 
 from ocf_data_sampler.config.load import load_yaml_configuration
-from ocf_data_sampler.config.model import Configuration, InputData
+from ocf_data_sampler.config.model import PVNetDataConfig
 from ocf_data_sampler.config.save import save_yaml_configuration
diff --git a/src/ocf_data_sampler/config/load.py b/src/ocf_data_sampler/config/load.py
index 150262e4..8fa08f39 100644
--- a/src/ocf_data_sampler/config/load.py
+++ b/src/ocf_data_sampler/config/load.py
@@ -3,10 +3,10 @@
 import fsspec
 from pyaml_env import parse_config
 
-from ocf_data_sampler.config.model import Configuration
+from ocf_data_sampler.config.model import PVNetDataConfig
 
 
-def load_yaml_configuration(filename: str) -> Configuration:
+def load_yaml_configuration(filename: str) -> PVNetDataConfig:
     """Load a yaml file which has a configuration in it.
 
     Args:
@@ -19,4 +19,4 @@ def load_yaml_configuration(filename: str) -> Configuration:
     with fsspec.open(filename, mode="r") as stream:
         configuration = parse_config(data=stream)
 
-    return Configuration(**configuration)
+    return PVNetDataConfig(**configuration)
diff --git a/src/ocf_data_sampler/config/model.py b/src/ocf_data_sampler/config/model.py
index f8c674ee..39611182 100644
--- a/src/ocf_data_sampler/config/model.py
+++ b/src/ocf_data_sampler/config/model.py
@@ -1,8 +1,4 @@
-"""Configuration model for the dataset.
-
-Absolute or relative zarr filepath(s).
-Prefix with a protocol like s3:// to read from alternative filesystems.
-"""
+"""Configuration model for the PVNet dataset."""
 
 from collections.abc import Iterator
 from typing import Literal
@@ -20,16 +16,6 @@ class Base(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
 
-class General(Base):
-    """General pydantic model."""
-
-    name: str = Field("example", description="The name of this configuration file")
-    description: str = Field(
-        "example configuration",
-        description="Description of this configuration file",
-    )
-
-
 class TimeWindowMixin(Base):
     """Mixin class, to add interval start, end and resolution minutes."""
 
@@ -72,7 +58,16 @@ def validate_intervals(self) -> "TimeWindowMixin":
         return self
 
 
-class DropoutMixin(Base):
+class FillValueMixin(Base):
+    """Mixin class, to add a value used for filling missing data."""
+
+    dropout_fill_value: float = Field(
+        default=0.0,
+        description="The value used to fill in dropped out data or any missing values."
+    )
+
+
+class DropoutMixin(FillValueMixin):
     """Mixin class, to add dropout minutes."""
 
     dropout_timedeltas_minutes: list[int] = Field(
@@ -87,11 +82,6 @@ class DropoutMixin(Base):
         "floats (probability that dropout of the corresponding timedelta is applied)",
     )
 
-    dropout_fill_value: float = Field(
-        default=0.0,
-        description="The value used to fill in dropped out data or any missing values."
-    )
-
     @field_validator("dropout_timedeltas_minutes")
     def dropout_timedeltas_minutes_negative(cls, v: list[int]) -> list[int]:
         """Validate 'dropout_timedeltas_minutes'."""
@@ -235,7 +225,7 @@ class NWP(TimeWindowMixin, DropoutMixin, SpatialWindowMixin, NormalisationConsta
 
     provider: str = Field(..., description="The provider of the NWP data")
 
-    accum_channels: list[str] = Field([], description="the nwp channels which need to be diffed")
+    accum_channels: list[str] = Field([], description="The NWP channels which need to be diffed")
 
     max_staleness_minutes: int | None = Field(
         None,
@@ -319,8 +309,58 @@ def items(self) -> Iterator[tuple[str, NWP]]:
         return self.root.items()
 
 
-class Generation(TimeWindowMixin, DropoutMixin):
-    """Generation configuration model."""
+class GenerationWindow(Base):
+    """Mixin class, to add interval start and end minutes for a generation window.
+
+    Unlike `TimeWindowMixin`, the temporal resolution is not included here - it belongs to the
+    shared generation data source (`Generation.time_resolution_minutes`), not to an individual
+    window over it.
+    """
+
+    interval_start_minutes: int = Field(
+        ...,
+        description="Data interval starts at `t0 + interval_start_minutes`",
+    )
+
+    interval_end_minutes: int = Field(
+        ...,
+        description="Data interval ends at `t0 + interval_end_minutes`",
+    )
+
+    @model_validator(mode="after")
+    def validate_interval_order(self) -> "GenerationWindow":
+        """Validator for time interval fields."""
+        start = self.interval_start_minutes
+        end = self.interval_end_minutes
+        if start > end:
+            raise ValueError(
+                f"interval_start_minutes ({start}) must be <= interval_end_minutes ({end})",
+            )
+        return self
+
+
+class GenerationInputWindow(GenerationWindow, DropoutMixin):
+    """Generation input window configuration model, used for `Generation.input`.
+
+    Extends `GenerationWindow` with dropout configuration, since only the input window (not the
+    prediction target) should ever be randomly masked out.
+    """
+
+
+class GenerationTargetWindow(GenerationWindow, FillValueMixin):
+    """Generation target window configuration model, used for `Generation.target`."""
+
+
+class Generation(Base):
+    """Generation configuration model.
+
+    Bundles the shared generation data source (`zarr_path`, `time_resolution_minutes`) with its
+    `input` and `target` windows - two independently configurable time windows over the same
+    underlying data. `time_resolution_minutes` describes generation's own native data cadence
+    (used for gap detection and windowed slicing of generation's own data) - it is independent
+    of `SamplingGrid.t0_resolution_minutes`, which is the cadence t0 candidates are enumerated
+    at and may legitimately differ (e.g. generation stored every 5 minutes, sampled every 30).
+    """
 
     zarr_path: str = Field(
         ...,
@@ -328,11 +368,69 @@ class Generation(TimeWindowMixin, DropoutMixin):
         "to read from alternative filesystems.",
     )
 
+    time_resolution_minutes: int = Field(
+        ...,
+        gt=0,
+        description="The temporal resolution of the generation data in minutes",
+    )
+
+    input: GenerationInputWindow | None = None
+    target: GenerationTargetWindow | None = None
+
+    @model_validator(mode="after")
+    def validate_windows(self) -> "Generation":
+        """Validate the input/target windows are set and divisible by the shared resolution."""
+        if self.input is None and self.target is None:
+            raise ValueError(
+                "At least one of `generation.input` or `generation.target` must be configured",
+            )
+
+        for name, window in (("input", self.input), ("target", self.target)):
+            if window is None:
+                continue
+            for bound_name, bound in (
+                ("interval_start_minutes", window.interval_start_minutes),
+                ("interval_end_minutes", window.interval_end_minutes),
+            ):
+                if bound % self.time_resolution_minutes != 0:
+                    raise ValueError(
+                        f"generation.{name}.{bound_name} ({bound}) must be divisible by "
+                        f"generation.time_resolution_minutes ({self.time_resolution_minutes})",
+                    )
+        return self
+
+
+class SamplingGrid(Base):
+    """Configuration for the (location, time) grid that t0 times are sampled from.
+
+    `locations_zarr_path` points to the locations metadata (location IDs and their
+    coordinates) - see `ocf_data_sampler.load.locations.open_locations`.
+    `t0_resolution_minutes` is the cadence t0 candidates are enumerated at, needed to compute
+    valid t0 times regardless of which other input sources are configured - it is not any one
+    source's own native data resolution (see `Generation.time_resolution_minutes` for that).
+    """
+
+    locations_zarr_path: str = Field(
+        ...,
+        description="Absolute or relative zarr filepath to the locations metadata. Prefix with "
+        "a protocol like s3:// to read from alternative filesystems.",
+    )
+
+    t0_resolution_minutes: int = Field(
+        ...,
+        gt=0,
+        description="The resolution of the t0 sampling grid, in minutes.",
+    )
+
 
 class SolarPosition(TimeWindowMixin):
     """Solar position configuration model."""
 
 
+class DatetimeEncoding(TimeWindowMixin):
+    """Datetime encoding configuration model."""
+
+
 _embedding_type = list[tuple[str, Literal["cyclic", "linear"]]]
 class T0Embedding(Base):
     """Configuration for the t0 time embedding."""
@@ -374,18 +472,13 @@ def validate_embeddings(cls, embeddings: _embedding_type) -> _embedding_type:
         return embeddings
 
 
-class InputData(Base):
-    """Input data model."""
+class PVNetDataConfig(Base):
+    """Configuration model for the PVNet dataset."""
 
+    sampling_grid: SamplingGrid
     satellite: Satellite | None = None
     nwp: MultiNWP | None = None
     generation: Generation | None = None
     solar_position: SolarPosition | None = None
+    datetime_encoding: DatetimeEncoding | None = None
     t0_embedding: T0Embedding | None = None
-
-
-class Configuration(Base):
-    """Configuration model for the dataset."""
-
-    general: General = General()
-    input_data: InputData = InputData()
diff --git a/src/ocf_data_sampler/config/save.py b/src/ocf_data_sampler/config/save.py
index 3bb15869..57bd96dc 100644
--- a/src/ocf_data_sampler/config/save.py
+++ b/src/ocf_data_sampler/config/save.py
@@ -10,14 +10,14 @@
 import fsspec
 import yaml
 
-from ocf_data_sampler.config.model import Configuration
+from ocf_data_sampler.config.model import PVNetDataConfig
 
 
-def save_yaml_configuration(configuration: Configuration, filename: str) -> None:
+def save_yaml_configuration(configuration: PVNetDataConfig, filename: str) -> None:
     """Save a configuration object to a YAML file.
 
     Args:
-        configuration: Configuration object containing the settings to save
+        configuration: PVNetDataConfig object containing the settings to save
         filename: Destination path for the YAML file. Can be a local path or
                  cloud storage URL (e.g., 'gs://', 's3://'). For local paths,
                  absolute paths are recommended.
diff --git a/src/ocf_data_sampler/datasets/pvnet/dataset.py b/src/ocf_data_sampler/datasets/pvnet/dataset.py
index c1f3dfed..27895212 100644
--- a/src/ocf_data_sampler/datasets/pvnet/dataset.py
+++ b/src/ocf_data_sampler/datasets/pvnet/dataset.py
@@ -12,15 +12,13 @@
 from ocf_data_sampler.common.lightarray import LightDataArray
 from ocf_data_sampler.common.time_utils import date_range, get_posix_timestamp, minutes
 from ocf_data_sampler.config.load import load_yaml_configuration
-from ocf_data_sampler.config.model import Configuration
+from ocf_data_sampler.config.model import PVNetDataConfig
 from ocf_data_sampler.datasets.cache import PickleCacheMixin
 from ocf_data_sampler.datasets.pvnet.loading import get_dataset_dict
 from ocf_data_sampler.datasets.pvnet.materialise import load_data_dict
 from ocf_data_sampler.datasets.pvnet.preprocess import (
-    apply_dropout_to_datasets,
     config_normalization_values_to_dicts,
-    diff_nwp_data,
-    fill_nans_in_dataset_dicts,
+    preprocess_dataset_dict,
 )
 from ocf_data_sampler.datasets.pvnet.sample import (
     convert_to_numpy_sample,
@@ -35,6 +33,7 @@
 from ocf_data_sampler.datasets.pvnet.types import NumpySample, SourceDict, TensorBatch
 from ocf_data_sampler.datasets.pvnet.valid_t0s import find_valid_time_periods
 from ocf_data_sampler.features.time_encodings import encode_datetimes
+from ocf_data_sampler.load import open_locations
 from ocf_data_sampler.select import (
     fill_time_periods,
     find_contiguous_t0_periods,
@@ -46,21 +45,23 @@
 
 
 
-def get_locations(generation_data: xr.DataArray) -> list[Location]:
-    """Get list of locations of all locations.
+def get_locations(zarr_path: str) -> list[Location]:
+    """Load the locations metadata and build the list of all locations.
 
     Args:
-        generation_data: xarray dataarray of generation data with location info
+        zarr_path: Path to the locations zarr data
     """
+    locations_data = open_locations(zarr_path)
+
     locations = []
-    location_ids = generation_data["location_id"].values
+    location_ids = locations_data["location_id"].values
 
     for location_id in location_ids:
-        gen_data = generation_data.sel(location_id=location_id)
+        loc_data = locations_data.sel(location_id=location_id)
         locations.append(
             Location(
-                x=gen_data["longitude"].values,
-                y=gen_data["latitude"].values,
+                x=loc_data["longitude"].values,
+                y=loc_data["latitude"].values,
                 coord_system="lon_lat",
                 id=int(location_id),
             ),
@@ -161,6 +162,69 @@ def add_alternate_coordinate_projections(
     return locations
 
 
+def build_numpy_sample(
+    dataset_dict: SourceDict,
+    t0: np.datetime64,
+    location: Location,
+    config: PVNetDataConfig,
+    include_extra_metadata: bool = False,
+) -> NumpySample:
+    """Convert data to numpy arrays and add auxiliary features.
+
+    Note: the data in `dataset_dict` is expected to already be preprocessed - see
+    `preprocess_dataset_dict`.
+
+    Args:
+        dataset_dict: Dictionary of xarray datasets
+        t0: init-time for sample
+        location: location of the sample
+        config: PVNetDataConfig object
+        include_extra_metadata: Whether to add additional non-essential metadata to the sample
+    """
+    # Convert all xarray modalities to a single NumpySample
+    sample = convert_to_numpy_sample(dataset_dict, include_extra_metadata)
+
+    sample["location_id"] = location.id
+    lon, lat = location.in_coord_system("lon_lat")
+
+    if include_extra_metadata:
+        sample["location_longitude"] = lon
+        sample["location_latitude"] = lat
+
+    # Add t0 embedding if configured
+    if config.t0_embedding is not None:
+        sample.update(
+            make_t0_encoding_numpy_sample(t0, config.t0_embedding.embeddings),
+        )
+
+    # Add datetime encodings if configured
+    if config.datetime_encoding is not None:
+        dt_config = config.datetime_encoding
+
+        datetimes = date_range(
+            t0 + minutes(dt_config.interval_start_minutes),
+            t0 + minutes(dt_config.interval_end_minutes),
+            freq=minutes(dt_config.time_resolution_minutes),
+        )
+        sample.update(encode_datetimes(datetimes=datetimes))
+
+    # Add solar position if configured
+    if config.solar_position is not None:
+        solar_config = config.solar_position
+
+        # Create datetime range for solar position calculation
+        datetimes = date_range(
+            t0 + minutes(solar_config.interval_start_minutes),
+            t0 + minutes(solar_config.interval_end_minutes),
+            freq=minutes(solar_config.time_resolution_minutes),
+        )
+
+        sample.update(make_sun_position_numpy_sample(datetimes, lon=lon, lat=lat))
+
+    sample["t0"] = get_posix_timestamp(t0)
+
+    return sample
+
 
 class AbstractPVNetDataset(PickleCacheMixin, Dataset):
     """Abstract class for PVNet datasets."""
@@ -191,10 +255,25 @@ def __init__(
 
         config = load_yaml_configuration(config_filename)
 
-        datasets_dict = get_dataset_dict(config.input_data)
+        locations = get_locations(config.sampling_grid.locations_zarr_path)
 
-        # Check if generation data has nans
-        self.complete_generation = not datasets_dict["generation"].isnull().any()
+        datasets_dict = get_dataset_dict(config)
+
+        if "generation" in datasets_dict:
+            location_ids = [loc.id for loc in locations]
+            missing = np.setdiff1d(location_ids, datasets_dict["generation"]["location_id"].values)
+            if len(missing) > 0:
+                raise ValueError(f"Generation data is missing for location IDs: {missing}")
+
+            # Restrict to the catalog's locations - generation may have extra ids that aren't
+            # real samplable points (e.g. a summation-model placeholder).
+            datasets_dict["generation"] = datasets_dict["generation"].sel(location_id=location_ids)
+
+        # Check if generation data has nans. If generation isn't configured at all, there's no
+        # per-location data availability to consider, so a single global t0 grid still applies.
+        self.complete_generation = (
+            "generation" not in datasets_dict or not datasets_dict["generation"].isnull().any()
+        )
 
         if self.complete_generation:
             valid_t0_times = self.find_valid_t0_times(datasets_dict, config)
@@ -219,9 +298,6 @@ def __init__(
 
             self.valid_t0_and_location_ids = valid_t0_and_location_ids
 
-        # Construct list of locations to sample from
-        locations = get_locations(generation_data=datasets_dict["generation"])
-
         self.locations = add_alternate_coordinate_projections(locations, datasets_dict)
 
         self.config = config
@@ -232,12 +308,6 @@ def __init__(
         else:
             self.datasets_dict = xarray_to_lightarray_dict(datasets_dict)
 
-        # Assign t0 idx value
-        self.t0_idx = (
-            -config.input_data.generation.interval_start_minutes
-            // config.input_data.generation.time_resolution_minutes
-        )
-
         # Extract the normalisation values from the config for faster access
         mean_dict, std_dict, clip_min_dict, clip_max_dict = (
             config_normalization_values_to_dicts(config)
@@ -262,113 +332,30 @@ def _sanitise_index(self, idx: int) -> int:
 
         return index
 
-    def process_and_combine_datasets(
-        self,
-        dataset_dict: SourceDict,
-        t0: np.datetime64,
-        location: Location,
-    ) -> NumpySample:
-        """Normalise and convert data to numpy arrays.
-
-        Args:
-            dataset_dict: Dictionary of xarray datasets
-            t0: init-time for sample
-            location: location of the sample
-        """
-        # Normalise NWP
-        if "nwp" in dataset_dict:
-            for nwp_key, da_nwp in dataset_dict["nwp"].items():
-                channel_means = self.mean_dict["nwp"][nwp_key]
-                channel_stds = self.std_dict["nwp"][nwp_key]
-                channel_mins = self.clip_min_dict["nwp"][nwp_key]
-                channel_maxs = self.clip_max_dict["nwp"][nwp_key]
-                dataset_dict["nwp"][nwp_key].data = (
-                    (da_nwp.data.clip(channel_mins, channel_maxs) - channel_means)
-                    / channel_stds
-                )
-
-        # Normalise satellite
-        if "sat" in dataset_dict:
-            channel_means = self.mean_dict["sat"]
-            channel_stds = self.std_dict["sat"]
-            channel_mins = self.clip_min_dict["sat"]
-            channel_maxs = self.clip_max_dict["sat"]
-            dataset_dict["sat"].data = (
-                (dataset_dict["sat"].data.clip(channel_mins, channel_maxs) - channel_means)
-                / channel_stds
-            )
-
-        # Fill NaNs
-        dataset_dict = fill_nans_in_dataset_dicts(dataset_dict, config=self.config)
-
-        # Convert all xarray modalities to a single NumpySample
-        sample = convert_to_numpy_sample(dataset_dict, self.t0_idx, self.include_extra_metadata)
-
-        # Add location metadata not present on the DataArray
-        if "generation" in dataset_dict:
-            sample["location_id"] = location.id
-
-        # Add datetime encodings over the full generation time range
-        generation_config = self.config.input_data.generation
-        datetimes = date_range(
-            t0 + minutes(generation_config.interval_start_minutes),
-            t0 + minutes(generation_config.interval_end_minutes),
-            freq=minutes(generation_config.time_resolution_minutes),
-        )
-        sample.update(encode_datetimes(datetimes=datetimes))
-
-        # Add t0 embedding if configured
-        if self.config.input_data.t0_embedding is not None:
-            sample.update(
-                make_t0_encoding_numpy_sample(t0, self.config.input_data.t0_embedding.embeddings),
-            )
-
-        # Add solar position if configured
-        if self.config.input_data.solar_position is not None:
-            solar_config = self.config.input_data.solar_position
-
-            # Create datetime range for solar position calculation
-            datetimes = date_range(
-                t0 + minutes(solar_config.interval_start_minutes),
-                t0 + minutes(solar_config.interval_end_minutes),
-                freq=minutes(solar_config.time_resolution_minutes),
-            )
-            sample.update(
-                make_sun_position_numpy_sample(
-                    datetimes,
-                    dataset_dict["generation"]["longitude"].values,
-                    dataset_dict["generation"]["latitude"].values,
-                ),
-            )
-
-        sample["t0"] = get_posix_timestamp(t0)
-
-        return sample
-
     @staticmethod
     def find_valid_t0_times(
         datasets_dict: SourceDict,
-        config: Configuration,
+        config: PVNetDataConfig,
     ) -> NDArray[np.datetime64]:
         """Find the t0 times where all of the requested input data is available.
 
         Args:
             datasets_dict: A dictionary of input datasets
-            config: Configuration file
+            config: PVNetDataConfig file
         """
         valid_time_periods = find_valid_time_periods(datasets_dict, config)
 
         # Fill out the contiguous time periods to get the t0 times
         valid_t0_times = fill_time_periods(
             valid_time_periods,
-            freq=minutes(config.input_data.generation.time_resolution_minutes),
+            freq=minutes(config.sampling_grid.t0_resolution_minutes),
         )
         return valid_t0_times
 
     @staticmethod
     def find_valid_t0_and_location_ids(
         datasets_dict: SourceDict,
-        config: Configuration,
+        config: PVNetDataConfig,
     ) -> pd.DataFrame:
         """Find the t0 times where all of the requested input data is available for each location.
 
@@ -378,7 +365,7 @@ def find_valid_t0_and_location_ids(
 
         Args:
             datasets_dict: A dictionary of input datasets
-            config: Configuration file
+            config: PVNetDataConfig file
         """
         # Get valid time period for nwp and satellite
         datasets_without_generation = {k: v for k, v in datasets_dict.items() if k != "generation"}
@@ -387,28 +374,34 @@ def find_valid_t0_and_location_ids(
         # Loop over each location in system id and obtain valid periods
         generations = datasets_dict["generation"]
         location_ids = generations.location_id.values
-        generation_config = config.input_data.generation
+        generation_config = config.generation
+        generation_windows = [
+            w for w in (generation_config.input, generation_config.target) if w is not None
+        ]
         valid_t0_and_location_ids = []
         for location_id in location_ids:
             generation = generations.sel(location_id=location_id)
             # Drop NaN values
             generation = generation.dropna(dim="time_utc")
 
-            # Obtain valid time periods for this location
-            time_periods = find_contiguous_t0_periods(
-                generation["time_utc"].values,
-                time_resolution=minutes(generation_config.time_resolution_minutes),
-                interval_start=minutes(generation_config.interval_start_minutes),
-                interval_end=minutes(generation_config.interval_end_minutes),
-            )
+            # Obtain valid time periods for this location, for each configured window
+            time_periods_per_window = [
+                find_contiguous_t0_periods(
+                    generation["time_utc"].values,
+                    time_resolution=minutes(generation_config.time_resolution_minutes),
+                    interval_start=minutes(window_config.interval_start_minutes),
+                    interval_end=minutes(window_config.interval_end_minutes),
+                )
+                for window_config in generation_windows
+            ]
             valid_time_periods_per_location = intersect_time_periods(
-                [valid_time_periods, time_periods],
+                [valid_time_periods, *time_periods_per_window],
             )
 
             # Fill out contiguous time periods to get t0 times
             valid_t0_times_per_location = fill_time_periods(
                 valid_time_periods_per_location,
-                freq=minutes(generation_config.time_resolution_minutes),
+                freq=minutes(config.sampling_grid.t0_resolution_minutes),
             )
 
             valid_t0_per_location = pd.DataFrame(index=valid_t0_times_per_location)
@@ -451,10 +444,13 @@ def _get_sample(self, t0: np.datetime64, location: Location) -> NumpySample:
         sample_dict = slice_datasets_by_space(self.datasets_dict, location, self.config)
         sample_dict = slice_datasets_by_time(sample_dict, t0, self.config)
         sample_dict = load_data_dict(sample_dict)
-        # Apply dropout to the data sources in-place
-        apply_dropout_to_datasets(sample_dict, t0, self.config)
-        sample_dict = diff_nwp_data(sample_dict, self.config)
-        return self.process_and_combine_datasets(sample_dict, t0, location)
+        sample_dict = preprocess_dataset_dict(
+            sample_dict, t0, self.config,
+            self.mean_dict, self.std_dict, self.clip_min_dict, self.clip_max_dict,
+        )
+        return build_numpy_sample(
+            sample_dict, t0, location, self.config, self.include_extra_metadata,
+        )
 
     @override
     def __getitem__(self, idx: int) -> NumpySample:
@@ -549,19 +545,19 @@ def _get_sample(self, t0: np.datetime64) -> TensorBatch:
         # Slice by time then load to avoid loading the data multiple times from disk
         sample_dict = slice_datasets_by_time(self.datasets_dict, t0, self.config)
         sample_dict = load_data_dict(sample_dict)
-        # Apply dropout to the data sources in-place
-        apply_dropout_to_datasets(sample_dict, t0, self.config)
-        sample_dict = diff_nwp_data(sample_dict, self.config)
+        # Preprocessing is location-independent, so do it once before slicing per-location below
+        sample_dict = preprocess_dataset_dict(
+            sample_dict, t0, self.config,
+            self.mean_dict, self.std_dict, self.clip_min_dict, self.clip_max_dict,
+        )
 
         samples = []
 
         # Prepare sample for each location
         for location in self.locations:
             sliced_sample_dict = slice_datasets_by_space(sample_dict, location, self.config)
-            numpy_sample = self.process_and_combine_datasets(
-                sliced_sample_dict,
-                t0,
-                location,
+            numpy_sample = build_numpy_sample(
+                sliced_sample_dict, t0, location, self.config, self.include_extra_metadata,
             )
             samples.append(numpy_sample)
 
diff --git a/src/ocf_data_sampler/datasets/pvnet/loading.py b/src/ocf_data_sampler/datasets/pvnet/loading.py
index ef3d359e..ecd63f84 100644
--- a/src/ocf_data_sampler/datasets/pvnet/loading.py
+++ b/src/ocf_data_sampler/datasets/pvnet/loading.py
@@ -1,41 +1,34 @@
 """Loads all data sources."""
 
-import logging
-
 import xarray as xr
 
-from ocf_data_sampler.config import InputData
+from ocf_data_sampler.config import PVNetDataConfig
 from ocf_data_sampler.datasets.pvnet.types import SourceDict
 from ocf_data_sampler.load import open_generation, open_nwp, open_sat_data
 
-logger = logging.getLogger(__name__)
 
+def get_dataset_dict(config: PVNetDataConfig) -> SourceDict[xr.DataArray]:
+    """Construct dictionary of all of the per-sample input data sources.
 
-def get_dataset_dict(input_config: InputData) -> SourceDict[xr.DataArray]:
-    """Construct dictionary of all of the input data sources.
+    Locations metadata is not included here - unlike generation/nwp/satellite it isn't a
+    per-sample source, so it's loaded separately by the caller - see
+    `ocf_data_sampler.load.locations.open_locations`.
 
     Args:
-        input_config: InputData configuration object
+        config: PVNetDataConfig configuration object
     """
     datasets_dict = {}
 
-    # Load generation data unless the path is None
-    if input_config.generation and input_config.generation.zarr_path:
-        da_generation = open_generation(zarr_path=input_config.generation.zarr_path)
-
-        # Remove location_id 0 if more than one location present
-        if len(da_generation["location_id"]) > 1 and 0 in da_generation["location_id"].values:
-            da_generation = da_generation.drop_sel(location_id=0)
-            logger.info(
-                "Id 0 has been filtered out, this is only used for summation models.",
-            )
-
-        datasets_dict["generation"] = da_generation
+    # Load generation data unless not configured. Any locations generation has that aren't in
+    # the locations catalog (e.g. an id used only for summation models) get filtered out later,
+    # once the catalog is available - see AbstractPVNetDataset.__init__.
+    if config.generation is not None:
+        datasets_dict["generation"] = open_generation(zarr_path=config.generation.zarr_path)
 
     # Load NWP data if in config
-    if input_config.nwp:
+    if config.nwp:
         datasets_dict["nwp"] = {}
-        for nwp_source, nwp_config in input_config.nwp.items():
+        for nwp_source, nwp_config in config.nwp.items():
             da_nwp = open_nwp(zarr_path=nwp_config.zarr_path, provider=nwp_config.provider)
 
             da_nwp = da_nwp.sel(channel=list(nwp_config.channels))
@@ -43,12 +36,11 @@ def get_dataset_dict(input_config: InputData) -> SourceDict[xr.DataArray]:
             datasets_dict["nwp"][nwp_source] = da_nwp
 
     # Load satellite data if in config
-    if input_config.satellite:
-        sat_config = input_config.satellite
+    if config.satellite:
 
-        da_sat = open_sat_data(sat_config.zarr_path)
+        da_sat = open_sat_data(config.satellite.zarr_path)
 
-        da_sat = da_sat.sel(channel=list(sat_config.channels))
+        da_sat = da_sat.sel(channel=list(config.satellite.channels))
 
         datasets_dict["sat"] = da_sat
 
diff --git a/src/ocf_data_sampler/datasets/pvnet/preprocess.py b/src/ocf_data_sampler/datasets/pvnet/preprocess.py
index e7cf44c6..2759b7bc 100644
--- a/src/ocf_data_sampler/datasets/pvnet/preprocess.py
+++ b/src/ocf_data_sampler/datasets/pvnet/preprocess.py
@@ -4,14 +4,14 @@
 
 from ocf_data_sampler.common.time_utils import minutes
 from ocf_data_sampler.common.types import TArray
-from ocf_data_sampler.config.model import Configuration
+from ocf_data_sampler.config.model import PVNetDataConfig
 from ocf_data_sampler.datasets.pvnet.types import SourceDict
 from ocf_data_sampler.features.diff_channels import diff_channels
-from ocf_data_sampler.select.dropout import apply_history_dropout
+from ocf_data_sampler.select.dropout import apply_dropout
 
 
 def config_normalization_values_to_dicts(
-    config: Configuration,
+    config: PVNetDataConfig,
 ) -> tuple[dict[str, np.ndarray | dict[str, np.ndarray]]]:
     """Construct numpy arrays of mean, std, and clip values from the config normalisation constants.
 
@@ -29,15 +29,14 @@ def config_normalization_values_to_dicts(
     clip_min_dict = {}
     clip_max_dict = {}
 
-    if config.input_data.nwp is not None:
+    if config.nwp is not None:
 
         means_dict["nwp"] = {}
         stds_dict["nwp"] = {}
         clip_min_dict["nwp"] = {}
         clip_max_dict["nwp"] = {}
 
-        for nwp_key in config.input_data.nwp:
-            nwp_config = config.input_data.nwp[nwp_key]
+        for nwp_key, nwp_config in config.nwp.items():
 
             means_list = []
             stds_list = []
@@ -61,16 +60,15 @@ def config_normalization_values_to_dicts(
             clip_min_dict["nwp"][nwp_key] = np.array(clip_min_list)[None, :, None, None]
             clip_max_dict["nwp"][nwp_key] = np.array(clip_max_list)[None, :, None, None]
 
-    if config.input_data.satellite is not None:
-        sat_config = config.input_data.satellite
+    if config.satellite is not None:
 
         means_list = []
         stds_list = []
         clip_min_list = []
         clip_max_list = []
 
-        for channel in list(sat_config.channels):
-            norm_conf = sat_config.normalisation_constants[channel]
+        for channel in list(config.satellite.channels):
+            norm_conf = config.satellite.normalisation_constants[channel]
             means_list.append(norm_conf.mean)
             stds_list.append(norm_conf.std)
             clip_min_list.append(-np.inf if norm_conf.clip_min is None else norm_conf.clip_min)
@@ -85,16 +83,84 @@ def config_normalization_values_to_dicts(
     return means_dict, stds_dict, clip_min_dict, clip_max_dict
 
 
-def diff_nwp_data(dataset_dict: SourceDict, config: Configuration) -> SourceDict:
+def normalise_dataset_dicts(
+    dataset_dict: SourceDict,
+    mean_dict: dict[str, np.ndarray | dict[str, np.ndarray]],
+    std_dict: dict[str, np.ndarray | dict[str, np.ndarray]],
+    clip_min_dict: dict[str, np.ndarray | dict[str, np.ndarray]],
+    clip_max_dict: dict[str, np.ndarray | dict[str, np.ndarray]],
+) -> SourceDict:
+    """Normalise the NWP, satellite, and generation data in-place.
+
+    NWP and satellite are normalised using the per-channel mean/std/clip constants from config.
+    Generation is normalised differently: `generation_mw` is rescaled to a capacity factor by
+    dividing by `capacity_mwp`, which is time-varying, per-location data rather than a config
+    constant - so it can't use the same clip/mean/std path. `capacity_mwp` itself is left
+    unchanged, since it's exposed raw in the output sample.
+
+    Args:
+        dataset_dict: Dictionary of xarray datasets
+        mean_dict: Means, as constructed by `config_normalization_values_to_dicts`
+        std_dict: Standard deviations, as constructed by `config_normalization_values_to_dicts`
+        clip_min_dict: Clip minimums, as constructed by `config_normalization_values_to_dicts`
+        clip_max_dict: Clip maximums, as constructed by `config_normalization_values_to_dicts`
+    """
+    if "nwp" in dataset_dict:
+        for nwp_key, da_nwp in dataset_dict["nwp"].items():
+            channel_means = mean_dict["nwp"][nwp_key]
+            channel_stds = std_dict["nwp"][nwp_key]
+            channel_mins = clip_min_dict["nwp"][nwp_key]
+            channel_maxs = clip_max_dict["nwp"][nwp_key]
+            dataset_dict["nwp"][nwp_key].data = (
+                (da_nwp.data.clip(channel_mins, channel_maxs) - channel_means)
+                / channel_stds
+            )
+
+    if "sat" in dataset_dict:
+        channel_means = mean_dict["sat"]
+        channel_stds = std_dict["sat"]
+        channel_mins = clip_min_dict["sat"]
+        channel_maxs = clip_max_dict["sat"]
+        dataset_dict["sat"].data = (
+            (dataset_dict["sat"].data.clip(channel_mins, channel_maxs) - channel_means)
+            / channel_stds
+        )
+
+    for key in ("generation_input", "generation_target"):
+        if key not in dataset_dict:
+            continue
+
+        da = dataset_dict[key]
+        gen_idx = list(da["gen_param"].values).index("generation_mw")
+        cap_idx = list(da["gen_param"].values).index("capacity_mwp")
+
+        generation_values = da.isel(gen_param=gen_idx).values
+        capacity_values = da.isel(gen_param=cap_idx).values
+
+        # capacity_mwp is time-varying (per timestep, per location) - normalise element-wise
+        # rather than by a single scalar. Where capacity is 0 the ratio is undefined, so we
+        # emit NaN rather than silently switching units (raw MW) or dividing by zero - dropout
+        # fill (later in the pipeline) replaces it with a fixed value in normalised units.
+        da.data[..., gen_idx] = np.divide(
+            generation_values,
+            capacity_values,
+            out=np.full_like(generation_values, np.nan, dtype=float),
+            where=capacity_values != 0,
+        )
+
+    return dataset_dict
+
+
+def diff_nwp_data(dataset_dict: SourceDict, config: PVNetDataConfig) -> SourceDict:
     """Take the in-place diff of some channels of the NWP data.
 
     Args:
         dataset_dict: Dictionary of xarray datasets
-        config: Configuration object
+        config: PVNetDataConfig object
     """
     if "nwp" in dataset_dict:
         for nwp_key, da_nwp in dataset_dict["nwp"].items():
-            accum_channels = config.input_data.nwp[nwp_key].accum_channels
+            accum_channels = config.nwp[nwp_key].accum_channels
             if len(accum_channels)>0:
                 # diff_channels() is an in-place operation and modifies the input
                 dataset_dict["nwp"][nwp_key] = diff_channels(da_nwp, accum_channels)
@@ -104,61 +170,63 @@ def diff_nwp_data(dataset_dict: SourceDict, config: Configuration) -> SourceDict
 def apply_dropout_to_datasets(
     datasets_dict: SourceDict,
     t0: np.datetime64,
-    config: Configuration,
+    config: PVNetDataConfig,
 ) -> None:
     """Apply dropout in-placeto the dictionary of input data sources around a given t0 time.
 
     Args:
         datasets_dict: Dictionary of the input data sources
         t0: The init-time
-        config: Configuration object.
+        config: PVNetDataConfig object.
 
     Returns:
         None. The input datasets_dict is modified in place.
     """
     if "sat" in datasets_dict:
-        sat_config = config.input_data.satellite
-
-        apply_history_dropout(
+        apply_dropout(
             datasets_dict["sat"],
             t0,
-            dropout_timedeltas=minutes(sat_config.dropout_timedeltas_minutes),
-            dropout_frac=sat_config.dropout_fraction,
+            dropout_timedeltas=minutes(config.satellite.dropout_timedeltas_minutes),
+            dropout_frac=config.satellite.dropout_fraction,
         )
 
-    if "generation" in datasets_dict:
-        generation_config = config.input_data.generation
+    if "generation_input" in datasets_dict:
 
-        # Dropout on the past generation, but not the future generation
-        apply_history_dropout(
-            datasets_dict["generation"],
+        # capacity_mwp is dropped out along with generation_mw - if the input feed was stale for
+        # this timestep, we didn't know the capacity at that point either.
+        # generation_target is never dropped out - it's the prediction target, not an input.
+        apply_dropout(
+            datasets_dict["generation_input"],
             t0,
-            dropout_timedeltas=minutes(generation_config.dropout_timedeltas_minutes),
-            dropout_frac=generation_config.dropout_fraction,
+            dropout_timedeltas=minutes(config.generation.input.dropout_timedeltas_minutes),
+            dropout_frac=config.generation.input.dropout_fraction,
         )
 
     return
 
 
-def fill_nans_in_dataset_dicts(datasets_dict: SourceDict, config: Configuration) -> SourceDict:
+def fill_nans_in_dataset_dicts(datasets_dict: SourceDict, config: PVNetDataConfig) -> SourceDict:
     """Fills all NaN values in the dataarrays in-place.
 
     Args:
         datasets_dict: Dictionary of the input data sources
-        config: Configuration object.
+        config: PVNetDataConfig object.
     """
-    conf_in = config.input_data
-    if "generation" in datasets_dict:
-        datasets_dict["generation"] = fill_nans(
-            datasets_dict["generation"],
-            conf_in.generation.dropout_fill_value,
-        )
+    if config.generation is not None:
+        for key, window_config in (
+            ("generation_input", config.generation.input),
+            ("generation_target", config.generation.target),
+        ):
+            if key in datasets_dict:
+                datasets_dict[key] = fill_nans(
+                    datasets_dict[key], window_config.dropout_fill_value,
+                )
 
     if "sat" in datasets_dict:
-        datasets_dict["sat"] = fill_nans(datasets_dict["sat"], conf_in.satellite.dropout_fill_value)
+        datasets_dict["sat"] = fill_nans(datasets_dict["sat"], config.satellite.dropout_fill_value)
 
     if "nwp" in datasets_dict:
-        for nwp_key, nwp_config in config.input_data.nwp.items():
+        for nwp_key, nwp_config in config.nwp.items():
             datasets_dict["nwp"][nwp_key] = fill_nans(
                 datasets_dict["nwp"][nwp_key],
                 nwp_config.dropout_fill_value,
@@ -172,3 +240,42 @@ def fill_nans(da: TArray, fill_value: float) -> TArray:
     if np.isnan(da.data).any():
         da.data = np.nan_to_num(da.data, copy=True, nan=fill_value)
     return da
+
+
+def preprocess_dataset_dict(
+    dataset_dict: SourceDict,
+    t0: np.datetime64,
+    config: PVNetDataConfig,
+    mean_dict: dict[str, np.ndarray | dict[str, np.ndarray]],
+    std_dict: dict[str, np.ndarray | dict[str, np.ndarray]],
+    clip_min_dict: dict[str, np.ndarray | dict[str, np.ndarray]],
+    clip_max_dict: dict[str, np.ndarray | dict[str, np.ndarray]],
+) -> SourceDict:
+    """Diff, normalise, dropout, and fill NaNs in the dictionary of input data sources.
+
+    These steps are always applied in the order listed, since some steps depend on the output of
+    previous steps. For example,
+    - NWP channel differencing must be done before normalisation, since the diffed channels have
+      different statistics
+    - Dropout must be applied after normalisation, since the fill value is specified in normalised
+      units
+    - NaN filling must be done after dropout, since dropout introduces NaNs in the data
+
+    Note: `dataset_dict` is expected to already be loaded - see `load_data_dict`.
+
+    Args:
+        dataset_dict: Dictionary of xarray datasets
+        t0: The init-time
+        config: PVNetDataConfig object
+        mean_dict: Means, as constructed by `config_normalization_values_to_dicts`
+        std_dict: Standard deviations, as constructed by `config_normalization_values_to_dicts`
+        clip_min_dict: Clip minimums, as constructed by `config_normalization_values_to_dicts`
+        clip_max_dict: Clip maximums, as constructed by `config_normalization_values_to_dicts`
+    """
+    dataset_dict = diff_nwp_data(dataset_dict, config)
+    dataset_dict = normalise_dataset_dicts(
+        dataset_dict, mean_dict, std_dict, clip_min_dict, clip_max_dict,
+    )
+    apply_dropout_to_datasets(dataset_dict, t0, config)
+    dataset_dict = fill_nans_in_dataset_dicts(dataset_dict, config=config)
+    return dataset_dict
diff --git a/src/ocf_data_sampler/datasets/pvnet/sample.py b/src/ocf_data_sampler/datasets/pvnet/sample.py
index 9c6b33d5..524ddf5f 100644
--- a/src/ocf_data_sampler/datasets/pvnet/sample.py
+++ b/src/ocf_data_sampler/datasets/pvnet/sample.py
@@ -10,7 +10,6 @@
 
 def convert_to_numpy_sample(
     datasets_dict: SourceDict,
-    t0_idx: int,
     include_extra_metadata: bool = False,
 ) -> NumpySample:
     """Convert a dictionary of xarray objects to a NumpySample.
@@ -18,10 +17,10 @@ def convert_to_numpy_sample(
     Args:
         datasets_dict: Dictionary of xarray DataArrays, with same structure as used inside
             PVNetDataset classes. Expected keys are any of following:
-            - "generation": DataArray of generation data
+            - "generation_input": DataArray of generation data used as model input
+            - "generation_target": DataArray of generation data used as the prediction target
             - "sat": DataArray of satellite data
             - "nwp": dict of DataArrays by provider name (e.g. {"ukv": da, "ecmwf": da})
-        t0_idx: Index of t0 within generation
         include_extra_metadata: Whether to add additional non-essential metadata to the batch
 
     Returns:
@@ -29,36 +28,25 @@ def convert_to_numpy_sample(
     """
     numpy_sample: NumpySample = {}
 
-    if "generation" in datasets_dict:
-        da = datasets_dict["generation"]
+    for key in ("generation_input", "generation_target"):
+        if key not in datasets_dict:
+            continue
 
-        # Get the position index of the generation and capacities
+        da = datasets_dict[key]
+
+        # generation_mw has already been normalised in-place to a capacity factor so should be in
+        # range [0, 1]. capacity_mwp is still in MW
         gen_idx = list(da["gen_param"].values).index("generation_mw")
         cap_idx = list(da["gen_param"].values).index("capacity_mwp")
 
-        generation_values = da.isel(gen_param=gen_idx).values
-        capacity_value = da.isel(gen_param=cap_idx).values[0]
-
-        if capacity_value!=0:
-            generation_values = generation_values/capacity_value
-
         numpy_sample.update(
             {
-                "generation": generation_values,
-                "capacity_mwp": capacity_value,
-                "generation_t0_idx": int(t0_idx),
-                "generation_time_utc": da["time_utc"].values.astype(float),
+                key: da.isel(gen_param=gen_idx).values,
+                f"{key}_capacity_mwp": da.isel(gen_param=cap_idx).values,
+                f"{key}_time_utc": da["time_utc"].values.astype(float),
             },
         )
 
-        if include_extra_metadata:
-            numpy_sample.update(
-                {
-                    "location_longitude": float(da["longitude"].values),
-                    "location_latitude": float(da["latitude"].values),
-                },
-            )
-
     if "sat" in datasets_dict:
         da = datasets_dict["sat"]
         numpy_sample.update({"satellite": da.values})
diff --git a/src/ocf_data_sampler/datasets/pvnet/slicing.py b/src/ocf_data_sampler/datasets/pvnet/slicing.py
index 0bf33029..fd170560 100644
--- a/src/ocf_data_sampler/datasets/pvnet/slicing.py
+++ b/src/ocf_data_sampler/datasets/pvnet/slicing.py
@@ -4,7 +4,7 @@
 
 from ocf_data_sampler.common.indexing import get_indices_in_sorted_unique
 from ocf_data_sampler.common.time_utils import minutes
-from ocf_data_sampler.config.model import Configuration
+from ocf_data_sampler.config.model import PVNetDataConfig
 from ocf_data_sampler.datasets.pvnet.types import SourceDict
 from ocf_data_sampler.select.spatial_slice import (
     select_spatial_slice_pixels,
@@ -17,29 +17,24 @@
 def slice_datasets_by_space(
     datasets_dict: SourceDict,
     location: Location,
-    config: Configuration,
+    config: PVNetDataConfig,
 ) -> SourceDict:
     """Slice the dictionary of input data sources around a given location.
 
     Args:
         datasets_dict: Dictionary of the input data sources
         location: The location to sample around
-        config: Configuration object.
+        config: PVNetDataConfig object.
 
     Returns:
         A dictionary of the sliced input data sources.
     """
-    if not set(datasets_dict.keys()).issubset({"nwp", "sat", "generation"}):
-        raise ValueError(
-            "'datasets_dict' should only contain keys 'nwp', 'sat', 'generation'",
-        )
-
     sliced_datasets_dict = {}
 
     if "nwp" in datasets_dict:
         sliced_datasets_dict["nwp"] = {}
 
-        for nwp_key, nwp_config in config.input_data.nwp.items():
+        for nwp_key, nwp_config in config.nwp.items():
             sliced_datasets_dict["nwp"][nwp_key] = select_spatial_slice_pixels(
                 datasets_dict["nwp"][nwp_key],
                 location,
@@ -48,21 +43,24 @@ def slice_datasets_by_space(
             )
 
     if "sat" in datasets_dict:
-        sat_config = config.input_data.satellite
-
         sliced_datasets_dict["sat"] = select_spatial_slice_pixels(
             datasets_dict["sat"],
             location,
-            height_pixels=sat_config.image_size_pixels_height,
-            width_pixels=sat_config.image_size_pixels_width,
+            height_pixels=config.satellite.image_size_pixels_height,
+            width_pixels=config.satellite.image_size_pixels_width,
         )
 
-    if "generation" in datasets_dict:
+    # Depending on whether this is called before or after time-slicing, the generation data is
+    # under a single "generation" key (raw, pre-split) or "generation_input"/"generation_target"
+    # (post-split) - slice whichever of these are present by location.
+    for key in ("generation", "generation_input", "generation_target"):
+        if key not in datasets_dict:
+            continue
 
-        location_ids = datasets_dict["generation"]["location_id"].values
+        location_ids = datasets_dict[key]["location_id"].values
         loc_index = get_indices_in_sorted_unique(location_ids, location.id)
 
-        sliced_datasets_dict["generation"] = datasets_dict["generation"].isel(location_id=loc_index)
+        sliced_datasets_dict[key] = datasets_dict[key].isel(location_id=loc_index)
 
     return sliced_datasets_dict
 
@@ -70,14 +68,14 @@ def slice_datasets_by_space(
 def reduce_spatial_extent_of_datasets(
     datasets_dict: SourceDict,
     locations: list[Location],
-    config: Configuration,
+    config: PVNetDataConfig,
 ) -> SourceDict:
     """Reduce the spatial extent of the datasets to only cover the locations.
 
     Args:
         datasets_dict: Dictionary of the input data sources
         locations: List of locations to reduce to
-        config: Configuration object
+        config: PVNetDataConfig object
 
     Returns:
         A dictionary of the reduced input data sources.
@@ -87,7 +85,7 @@ def reduce_spatial_extent_of_datasets(
     if "nwp" in datasets_dict:
         sliced_datasets_dict["nwp"] = {}
 
-        for nwp_key, nwp_config in config.input_data.nwp.items():
+        for nwp_key, nwp_config in config.nwp.items():
             sliced_datasets_dict["nwp"][nwp_key] = select_spatial_slice_pixels_multiple(
                 datasets_dict["nwp"][nwp_key],
                 locations,
@@ -97,7 +95,7 @@ def reduce_spatial_extent_of_datasets(
 
 
     if "sat" in datasets_dict:
-        sat_config = config.input_data.satellite
+        sat_config = config.satellite
 
         sliced_datasets_dict["sat"] = select_spatial_slice_pixels_multiple(
             datasets_dict["sat"],
@@ -115,14 +113,14 @@ def reduce_spatial_extent_of_datasets(
 def slice_datasets_by_time(
     datasets_dict: SourceDict,
     t0: np.datetime64,
-    config: Configuration,
+    config: PVNetDataConfig,
 ) -> SourceDict:
     """Slice the dictionary of input data sources around a given t0 time.
 
     Args:
         datasets_dict: Dictionary of the input data sources
         t0: The init-time
-        config: Configuration object.
+        config: PVNetDataConfig object.
 
     Returns:
         A dictionary of the sliced input data sources.
@@ -133,7 +131,7 @@ def slice_datasets_by_time(
         sliced_datasets_dict["nwp"] = {}
 
         for nwp_key, da_nwp in datasets_dict["nwp"].items():
-            nwp_config = config.input_data.nwp[nwp_key]
+            nwp_config = config.nwp[nwp_key]
 
             # Add a buffer if we need to diff some of the channels in time
             if len(nwp_config.accum_channels)>0:
@@ -155,7 +153,7 @@ def slice_datasets_by_time(
             )
 
     if "sat" in datasets_dict:
-        sat_config = config.input_data.satellite
+        sat_config = config.satellite
 
         sliced_datasets_dict["sat"] = select_time_slice(
             datasets_dict["sat"],
@@ -166,14 +164,20 @@ def slice_datasets_by_time(
         )
 
     if "generation" in datasets_dict:
-        generation_config = config.input_data.generation
-
-        sliced_datasets_dict["generation"] = select_time_slice(
-            datasets_dict["generation"],
-            t0,
-            time_resolution=minutes(generation_config.time_resolution_minutes),
-            interval_start=minutes(generation_config.interval_start_minutes),
-            interval_end=minutes(generation_config.interval_end_minutes),
-        )
+        generation_config = config.generation
+        for key, window_config in (
+            ("generation_input", generation_config.input),
+            ("generation_target", generation_config.target),
+        ):
+            if window_config is None:
+                continue
+
+            sliced_datasets_dict[key] = select_time_slice(
+                datasets_dict["generation"],
+                t0,
+                time_resolution=minutes(generation_config.time_resolution_minutes),
+                interval_start=minutes(window_config.interval_start_minutes),
+                interval_end=minutes(window_config.interval_end_minutes),
+            )
 
     return sliced_datasets_dict
diff --git a/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py b/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py
index 8514fe12..a824fc87 100644
--- a/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py
+++ b/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py
@@ -4,7 +4,7 @@
 import pandas as pd
 
 from ocf_data_sampler.common.time_utils import minutes
-from ocf_data_sampler.config.model import Configuration
+from ocf_data_sampler.config.model import PVNetDataConfig
 from ocf_data_sampler.datasets.pvnet.types import SourceDict
 from ocf_data_sampler.select.time_periods import (
     find_contiguous_t0_periods,
@@ -15,20 +15,20 @@
 
 def find_valid_time_periods(
     datasets_dict: SourceDict,
-    config: Configuration,
+    config: PVNetDataConfig,
 ) -> pd.DataFrame:
     """Find the t0 times where all of the requested input data is available.
 
     Args:
         datasets_dict: A dictionary of input datasets
-        config: Configuration file
+        config: PVNetDataConfig file
 
     Returns:
         A DataFrame containing the valid t0 time periods.
     """
     contiguous_time_periods: list[pd.DataFrame] = []
     if "nwp" in datasets_dict:
-        for nwp_key, nwp_config in config.input_data.nwp.items():
+        for nwp_key, nwp_config in config.nwp.items():
             da = datasets_dict["nwp"][nwp_key]
 
             # Extract the max extents of the forecast steps
@@ -68,13 +68,11 @@ def find_valid_time_periods(
             contiguous_time_periods.append(time_periods)
 
     if "sat" in datasets_dict:
-        sat_config = config.input_data.satellite
-
         time_periods = find_contiguous_t0_periods(
             datasets_dict["sat"]["time_utc"].values,
-            time_resolution=minutes(sat_config.time_resolution_minutes),
-            interval_start=minutes(sat_config.interval_start_minutes),
-            interval_end=minutes(sat_config.interval_end_minutes),
+            time_resolution=minutes(config.satellite.time_resolution_minutes),
+            interval_start=minutes(config.satellite.interval_start_minutes),
+            interval_end=minutes(config.satellite.interval_end_minutes),
         )
 
         contiguous_time_periods.append(time_periods)
@@ -83,19 +81,21 @@ def find_valid_time_periods(
             raise ValueError("No valid t0 periods found for satellite data")
 
     if "generation" in datasets_dict:
-        generation_config = config.input_data.generation
-
-        time_periods = find_contiguous_t0_periods(
-            datasets_dict["generation"]["time_utc"].values,
-            time_resolution=minutes(generation_config.time_resolution_minutes),
-            interval_start=minutes(generation_config.interval_start_minutes),
-            interval_end=minutes(generation_config.interval_end_minutes),
-        )
+        for window_config in (config.generation.input, config.generation.target):
+            if window_config is None:
+                continue
+
+            time_periods = find_contiguous_t0_periods(
+                datasets_dict["generation"]["time_utc"].values,
+                time_resolution=minutes(config.generation.time_resolution_minutes),
+                interval_start=minutes(window_config.interval_start_minutes),
+                interval_end=minutes(window_config.interval_end_minutes),
+            )
 
-        if len(time_periods) == 0:
-            raise ValueError("No valid t0 periods found for generation data")
+            if len(time_periods) == 0:
+                raise ValueError("No valid t0 periods found for generation data")
 
-        contiguous_time_periods.append(time_periods)
+            contiguous_time_periods.append(time_periods)
 
     # Find joint overlapping contiguous time periods
     valid_time_periods = intersect_time_periods(contiguous_time_periods)
diff --git a/src/ocf_data_sampler/load/__init__.py b/src/ocf_data_sampler/load/__init__.py
index f2fe9000..17bdcf17 100644
--- a/src/ocf_data_sampler/load/__init__.py
+++ b/src/ocf_data_sampler/load/__init__.py
@@ -1,3 +1,4 @@
 from ocf_data_sampler.load.generation import open_generation
+from ocf_data_sampler.load.locations import open_locations
 from ocf_data_sampler.load.nwp import open_nwp
 from ocf_data_sampler.load.satellite import open_sat_data
diff --git a/src/ocf_data_sampler/load/generation.py b/src/ocf_data_sampler/load/generation.py
index b55d2f52..0d7691dd 100644
--- a/src/ocf_data_sampler/load/generation.py
+++ b/src/ocf_data_sampler/load/generation.py
@@ -9,9 +9,6 @@
 Coordinates:
     time_utc (time_utc): The datetimes associated with each generation and capacity value
     location_id (location_id): The integer IDs of the locations
-    longitude (location_id): The longitudes of the locations
-    latitude (location_id): The latitudes of the locations
-
 """
 
 import numpy as np
@@ -43,8 +40,6 @@ def open_generation(zarr_path: str) -> xr.DataArray:
     coord_dtypes = {
         "time_utc": np.datetime64,
         "location_id": np.integer,
-        "longitude": np.number,
-        "latitude": np.number,
     }
     validate_coords(
         ds,
diff --git a/src/ocf_data_sampler/load/locations.py b/src/ocf_data_sampler/load/locations.py
new file mode 100644
index 00000000..d888fac5
--- /dev/null
+++ b/src/ocf_data_sampler/load/locations.py
@@ -0,0 +1,47 @@
+"""Functions for loading locations metadata.
+
+Locations data schema: a Zarr file with the following data variables and dimensions/coordinates:
+
+Dimensions: (location_id,)
+Data Variables:
+    longitude (location_id): The longitudes of the locations
+    latitude (location_id): The latitudes of the locations
+Coordinates:
+    location_id (location_id): The integer IDs of the locations
+"""
+
+import numpy as np
+import xarray as xr
+
+from ocf_data_sampler.common.indexing import assert_values_unique_increasing
+from ocf_data_sampler.load.conventions import validate_coords
+
+
+def open_locations(zarr_path: str) -> xr.Dataset:
+    """Open and eagerly load the locations metadata and validate its data types.
+
+    Args:
+        zarr_path: Path to the locations zarr data
+
+    Returns:
+        xr.Dataset: The opened locations metadata
+    """
+    ds = xr.open_zarr(zarr_path, chunks=None)
+
+    if set(ds.data_vars) != {"longitude", "latitude"}:
+        raise ValueError(
+            f"Locations data should have variables 'longitude' and 'latitude', "
+            f"but found {set(ds.data_vars)} instead."
+        )
+
+    validate_coords(ds, {"location_id": np.integer}, source="locations data")
+
+    for var in ("longitude", "latitude"):
+        if not np.issubdtype(ds[var].dtype, np.floating):
+            raise TypeError(f"{var} in locations data should be floating, not {ds[var].dtype}")
+
+    ds = ds.load()
+
+    assert_values_unique_increasing(ds["location_id"].values, "location_id")
+
+    return ds
diff --git a/src/ocf_data_sampler/select/dropout.py b/src/ocf_data_sampler/select/dropout.py
index 8972ce8f..5465eb3f 100644
--- a/src/ocf_data_sampler/select/dropout.py
+++ b/src/ocf_data_sampler/select/dropout.py
@@ -5,15 +5,16 @@
 from ocf_data_sampler.common.types import TArray
 
 
-def apply_history_dropout(
+def apply_dropout(
     da: TArray,
     t0: np.datetime64,
     dropout_timedeltas: list[np.timedelta64],
     dropout_frac: float | list[float],
 ) -> TArray:
-    """Apply in-place random dropout to the historical part of some sequence data.
+    """Apply in-place random dropout to some sequence data.
 
-    Dropped out data is replaced with NaNs.
+    A timedelta relative to t0 is randomly sampled, and all data after that point in time is
+    dropped out (replaced with NaNs).
 
     This helper requires a NumPy-backed DataArray-like object. It mutates the backing array in
     place, so it should be called after any lazy data has been materialised.
@@ -61,9 +62,11 @@ def apply_history_dropout(
         return da
     else:
         times = da["time_utc"].values
-        keep = (times <= t0 + timedelta_choice) | (times > t0)
+        keep = times <= t0 + timedelta_choice
         axis = da.dims.index("time_utc")
         mask = np.expand_dims(keep, tuple(i for i in range(da.data.ndim) if i != axis))
+        mask = np.broadcast_to(mask, da.data.shape)
+
         da.data = np.where(mask, da.data, np.nan)
         return da
 
diff --git a/tests/config/test_config.py b/tests/config/test_config.py
index bb60dad2..d0e31714 100644
--- a/tests/config/test_config.py
+++ b/tests/config/test_config.py
@@ -2,34 +2,37 @@
 from pydantic import ValidationError
 
 from ocf_data_sampler.config.load import load_yaml_configuration
-from ocf_data_sampler.config.model import Configuration
+from ocf_data_sampler.config.model import PVNetDataConfig
 
 
 def _load_config_and_provider(config_path):
     config = load_yaml_configuration(config_path)
-    provider = next(iter(config.input_data.nwp.root.keys()))
+    provider = next(iter(config.nwp.root.keys()))
     return config, provider
 
 
 def _validate_configuration(config):
     """Recreate config instance from dict to trigger validation."""
-    return Configuration(**config.model_dump())
+    return PVNetDataConfig(**config.model_dump())
+
+
+_MINIMAL_SAMPLING_GRID = {"locations_zarr_path": "locations.zarr", "t0_resolution_minutes": 30}
 
 
 def test_default_configuration():
-    """Test default pydantic class"""
-    _ = Configuration()
+    """Test default pydantic class - sampling_grid is the only required field"""
+    _ = PVNetDataConfig(sampling_grid=_MINIMAL_SAMPLING_GRID)
 
 
 def test_extra_field_error():
     """
     Check an extra parameters in config causes error
     """
-    configuration = Configuration()
+    configuration = PVNetDataConfig(sampling_grid=_MINIMAL_SAMPLING_GRID)
     configuration_dict = configuration.model_dump()
     configuration_dict["extra_field"] = "extra_value"
     with pytest.raises(ValidationError, match="Extra inputs are not permitted"):
-        Configuration(**configuration_dict)
+        PVNetDataConfig(**configuration_dict)
 
 
 def test_incorrect_interval_start_minutes(config_filename):
@@ -37,7 +40,7 @@ def test_incorrect_interval_start_minutes(config_filename):
     Check a history length not divisible by time resolution causes error
     """
     configuration, provider = _load_config_and_provider(config_filename)
-    configuration.input_data.nwp[provider].interval_start_minutes = -1111
+    configuration.nwp[provider].interval_start_minutes = -1111
     with pytest.raises(
         ValueError,
         match=r"interval_start_minutes \(-1111\) "
@@ -51,7 +54,7 @@ def test_incorrect_interval_end_minutes(config_filename):
     Check a forecast length not divisible by time resolution causes error
     """
     configuration, provider = _load_config_and_provider(config_filename)
-    configuration.input_data.nwp[provider].interval_end_minutes = 1111
+    configuration.nwp[provider].interval_end_minutes = 1111
     with pytest.raises(
         ValueError,
         match=r"interval_end_minutes \(1111\) "
@@ -65,7 +68,7 @@ def test_incorrect_nwp_provider(config_filename):
     Check an unexpected nwp provider causes error
     """
     configuration, provider = _load_config_and_provider(config_filename)
-    configuration.input_data.nwp[provider].provider = "unexpected_provider"
+    configuration.nwp[provider].provider = "unexpected_provider"
     with pytest.raises(ValidationError, match="Unknown NWP provider"):
         _validate_configuration(configuration)
 
@@ -73,11 +76,11 @@ def test_incorrect_nwp_provider(config_filename):
 def test_nwp_provider_is_canonicalized(config_filename):
     """NWP provider names are stored using their canonical lowercase spelling."""
     configuration, provider = _load_config_and_provider(config_filename)
-    configuration.input_data.nwp[provider].provider = "UKV"
+    configuration.nwp[provider].provider = "UKV"
 
     validated = _validate_configuration(configuration)
 
-    assert validated.input_data.nwp[provider].provider == "ukv"
+    assert validated.nwp[provider].provider == "ukv"
 
 
 def test_incorrect_dropout(config_filename):
@@ -87,12 +90,12 @@ def test_incorrect_dropout(config_filename):
     configuration, provider = _load_config_and_provider(config_filename)
 
     # Check that a positive number is not allowed
-    configuration.input_data.nwp[provider].dropout_timedeltas_minutes = [120]
+    configuration.nwp[provider].dropout_timedeltas_minutes = [120]
     with pytest.raises(Exception, match="Dropout timedeltas must be negative"):
         _validate_configuration(configuration)
 
     # Check that zero is allowed
-    configuration.input_data.nwp[provider].dropout_timedeltas_minutes = [0]
+    configuration.nwp[provider].dropout_timedeltas_minutes = [0]
     _validate_configuration(configuration)
 
 
@@ -102,23 +105,23 @@ def test_incorrect_dropout_fraction(config_filename):
     """
     configuration, provider = _load_config_and_provider(config_filename)
 
-    configuration.input_data.nwp[provider].dropout_fraction = 1.1
+    configuration.nwp[provider].dropout_fraction = 1.1
     with pytest.raises(ValidationError, match=r"Dropout fractions must be in range *"):
         _validate_configuration(configuration)
 
-    configuration.input_data.nwp[provider].dropout_fraction = -0.1
+    configuration.nwp[provider].dropout_fraction = -0.1
     with pytest.raises(ValidationError, match=r"Dropout fractions must be in range *"):
         _validate_configuration(configuration)
 
-    configuration.input_data.nwp[provider].dropout_fraction = [1.0, 0.1]
+    configuration.nwp[provider].dropout_fraction = [1.0, 0.1]
     with pytest.raises(ValidationError, match=r"The sum of dropout fractions must be in range *"):
         _validate_configuration(configuration)
 
-    configuration.input_data.nwp[provider].dropout_fraction = [-0.1, 1.1]
+    configuration.nwp[provider].dropout_fraction = [-0.1, 1.1]
     with pytest.raises(ValidationError, match=r"All dropout fractions must be in range *"):
         _validate_configuration(configuration)
 
-    configuration.input_data.nwp[provider].dropout_fraction = []
+    configuration.nwp[provider].dropout_fraction = []
     with pytest.raises(ValidationError, match="List cannot be empty"):
         _validate_configuration(configuration)
 
@@ -127,8 +130,8 @@ def test_dropout_fraction_list_length_matches_timedeltas(config_filename):
     """List dropout fractions must align with dropout timedeltas one-to-one."""
     configuration, provider = _load_config_and_provider(config_filename)
 
-    configuration.input_data.nwp[provider].dropout_timedeltas_minutes = [-60, -120]
-    configuration.input_data.nwp[provider].dropout_fraction = [0.5]
+    configuration.nwp[provider].dropout_timedeltas_minutes = [-60, -120]
+    configuration.nwp[provider].dropout_fraction = [0.5]
 
     with pytest.raises(
         ValidationError,
@@ -142,16 +145,16 @@ def test_inconsistent_dropout_use(config_filename):
     Check dropout fraction outside of range causes error
     """
     configuration = load_yaml_configuration(config_filename)
-    configuration.input_data.satellite.dropout_fraction = 1.0
-    configuration.input_data.satellite.dropout_timedeltas_minutes = []
+    configuration.satellite.dropout_fraction = 1.0
+    configuration.satellite.dropout_timedeltas_minutes = []
     with pytest.raises(
         ValueError,
         match="To dropout fraction > 0 requires a list of dropout timedeltas",
     ):
         _validate_configuration(configuration)
 
-    configuration.input_data.satellite.dropout_fraction = 0.0
-    configuration.input_data.satellite.dropout_timedeltas_minutes = [-120, -60]
+    configuration.satellite.dropout_fraction = 0.0
+    configuration.satellite.dropout_timedeltas_minutes = [-120, -60]
     with pytest.raises(
         ValueError,
         match="To use dropout timedeltas dropout fraction should be > 0",
@@ -165,15 +168,41 @@ def test_accum_channels_validation(config_filename):
 
     # Test invalid channel scenario
     invalid_config = config.model_copy(deep=True)
-    invalid_nwp = invalid_config.input_data.nwp.root[nwp_name]
+    invalid_nwp = invalid_config.nwp.root[nwp_name]
     invalid_nwp.accum_channels = ["invalid_channel"]
 
     # Verify exact error message
     expected_error = (
-        rf"input_data.nwp.{nwp_name}\n"
+        rf"nwp.{nwp_name}\n"
         fr"  Value error, NWP provider '{nwp_name}': all values in 'accum_channels' "
         r"should be present in 'channels'\. "
         r"Extra values found: {'invalid_channel'}.*"
     )
     with pytest.raises(ValidationError, match=expected_error):
         _validate_configuration(invalid_config)
+
+
+def test_generation_interval_divisibility_raises(config_filename):
+    """generation.input/target must be divisible by generation.time_resolution_minutes."""
+    configuration = load_yaml_configuration(config_filename)
+    configuration.generation.input.interval_start_minutes = -45
+
+    with pytest.raises(
+        ValueError,
+        match=r"generation\.input\.interval_start_minutes \(-45\) must be divisible by "
+        r"generation\.time_resolution_minutes \(30\)",
+    ):
+        _validate_configuration(configuration)
+
+
+def test_generation_requires_input_or_target(config_filename):
+    """At least one of generation.input or generation.target must be configured."""
+    configuration = load_yaml_configuration(config_filename)
+    configuration.generation.input = None
+    configuration.generation.target = None
+
+    with pytest.raises(
+        ValueError,
+        match=r"At least one of `generation\.input` or `generation\.target` must be configured",
+    ):
+        _validate_configuration(configuration)
diff --git a/tests/config/test_load.py b/tests/config/test_load.py
index 36f971a3..e2d1b1e9 100644
--- a/tests/config/test_load.py
+++ b/tests/config/test_load.py
@@ -1,6 +1,6 @@
 from ocf_data_sampler.config.load import load_yaml_configuration
-from ocf_data_sampler.config.model import Configuration
+from ocf_data_sampler.config.model import PVNetDataConfig
 
 
 def test_load_yaml_configuration(config_filename):
-    assert isinstance(load_yaml_configuration(config_filename), Configuration)
+    assert isinstance(load_yaml_configuration(config_filename), PVNetDataConfig)
diff --git a/tests/conftest.py b/tests/conftest.py
index 1f93d88e..db8bfa5a 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -283,18 +283,51 @@ def site_generation_zarr_path(session_tmp_path, ds_site_generation):
     yield save_zarr(ds_site_generation, session_tmp_path, "site_generation.zarr")
 
 
+def _locations_dataset_from_generation(ds_generation):
+    """Build a standalone locations metadata dataset from a generation dataset's location coords.
+
+    Excludes location_id 0 - in the generation fixtures that's a placeholder used only for
+    summation models, not a real samplable location, so a properly curated locations catalog
+    wouldn't list it even though generation does.
+    """
+    ds_generation = ds_generation.sel(
+        location_id=[loc_id for loc_id in ds_generation["location_id"].values if loc_id != 0],
+    )
+    return xr.Dataset(
+        data_vars={
+            "longitude": ("location_id", ds_generation["longitude"].values),
+            "latitude": ("location_id", ds_generation["latitude"].values),
+        },
+        coords={"location_id": ds_generation["location_id"].values},
+    )
+
+
+@pytest.fixture(scope="session")
+def locations_zarr_path(session_tmp_path, ds_generation):
+    ds_locations = _locations_dataset_from_generation(ds_generation)
+    yield save_zarr(ds_locations, session_tmp_path, "locations.zarr")
+
+
+@pytest.fixture(scope="session")
+def site_locations_zarr_path(session_tmp_path, ds_site_generation):
+    ds_locations = _locations_dataset_from_generation(ds_site_generation)
+    yield save_zarr(ds_locations, session_tmp_path, "site_locations.zarr")
+
+
 @pytest.fixture()
 def pvnet_config_filename(
     tmp_path,
     config_filename,
     nwp_ukv_zarr_path,
     generation_zarr_path,
+    locations_zarr_path,
     sat_zarr_path,
 ):
     config = load_yaml_configuration(config_filename)
-    config.input_data.nwp["ukv"].zarr_path = nwp_ukv_zarr_path
-    config.input_data.satellite.zarr_path = sat_zarr_path
-    config.input_data.generation.zarr_path = generation_zarr_path
+    config.nwp["ukv"].zarr_path = nwp_ukv_zarr_path
+    config.satellite.zarr_path = sat_zarr_path
+    config.generation.zarr_path = generation_zarr_path
+    config.sampling_grid.locations_zarr_path = locations_zarr_path
 
     path = tmp_path / "configuration.yaml"
     save_yaml_configuration(config, str(path))
@@ -307,12 +340,14 @@ def pvnet_site_config_filename(
     config_filename,
     nwp_ukv_zarr_path,
     site_generation_zarr_path,
+    site_locations_zarr_path,
     sat_zarr_path,
 ):
     config = load_yaml_configuration(config_filename)
-    config.input_data.nwp["ukv"].zarr_path = nwp_ukv_zarr_path
-    config.input_data.satellite.zarr_path = sat_zarr_path
-    config.input_data.generation.zarr_path = site_generation_zarr_path
+    config.nwp["ukv"].zarr_path = nwp_ukv_zarr_path
+    config.satellite.zarr_path = sat_zarr_path
+    config.generation.zarr_path = site_generation_zarr_path
+    config.sampling_grid.locations_zarr_path = site_locations_zarr_path
 
     path = session_tmp_path / "configuration.yaml"
     save_yaml_configuration(config, str(path))
diff --git a/tests/datasets/pvnet/test_dataset.py b/tests/datasets/pvnet/test_dataset.py
index 96d8ccce..b3d64f6c 100644
--- a/tests/datasets/pvnet/test_dataset.py
+++ b/tests/datasets/pvnet/test_dataset.py
@@ -24,21 +24,23 @@ def _pvnet_dataset_sample_check(sample, config, batch_dim = None):
     assert isinstance(sample, dict)
 
     # Specific keys should always be present
-    required_keys = ["nwp_ukv", "satellite", "generation", "t0", "t0_embedding"]
+    required_keys = [
+        "nwp_ukv", "satellite", "generation_input", "generation_target", "t0", "t0_embedding",
+    ]
     for key in required_keys:
         assert key in sample
 
     solar_keys = ["solar_azimuth", "solar_elevation"]
-    if config.input_data.solar_position is not None:
+    if config.solar_position is not None:
         # Test solar position keys are present when configured
         for key in solar_keys:
             assert key in sample, f"Solar position key {key} should be present in sample"
 
         # Get expected time steps from config
         expected_time_steps = (
-            config.input_data.solar_position.interval_end_minutes
-            - config.input_data.solar_position.interval_start_minutes
-        ) // config.input_data.solar_position.time_resolution_minutes + 1
+            config.solar_position.interval_end_minutes
+            - config.solar_position.interval_start_minutes
+        ) // config.solar_position.time_resolution_minutes + 1
 
         # Test solar angle shapes based on config
         assert sample["solar_azimuth"].shape == (*batch_dim, expected_time_steps)
@@ -48,16 +50,35 @@ def _pvnet_dataset_sample_check(sample, config, batch_dim = None):
         for key in solar_keys:
             assert key not in sample, f"Solar position key {key} should not be present"
 
+    datetime_encoding_keys = ["date_sin", "date_cos", "time_sin", "time_cos"]
+    if config.datetime_encoding is not None:
+        # Test datetime encoding keys are present when configured
+        for key in datetime_encoding_keys:
+            assert key in sample, f"Datetime encoding key {key} should be present in sample"
+
+        # Get expected time steps from config
+        expected_time_steps = (
+            config.datetime_encoding.interval_end_minutes
+            - config.datetime_encoding.interval_start_minutes
+        ) // config.datetime_encoding.time_resolution_minutes + 1
+
+        # Test datetime encoding shapes based on config
+        for key in datetime_encoding_keys:
+            assert sample[key].shape == (*batch_dim, expected_time_steps)
+    else:
+        # Assert that datetime encoding keys are not present
+        for key in datetime_encoding_keys:
+            assert key not in sample, f"Datetime encoding key {key} should not be present"
+
     # Check the shape of the data is correct
     # 30 minutes of 5 minute data (inclusive), one channel, 2x2 pixels
     assert sample["satellite"].shape == (*batch_dim, 7, 1, 2, 2)
     # 3 hours of 60 minute data (inclusive), one channel, 2x2 pixels
     assert sample["nwp_ukv"].shape == (*batch_dim, 4, 1, 2, 2)
-    # 3 hours of 30 minute data (inclusive)
-    assert sample["generation"].shape == (*batch_dim, 7)
-    # Datetime encoding keys same shape as the generation
-    for datetime_key in ["date_sin", "date_cos", "time_sin", "time_cos"]:
-        assert sample[datetime_key].shape == (*batch_dim, 7)
+    # generation_input: 1 hour of 30 minute data (inclusive) = 3 steps
+    # generation_target: 2 hours of 30 minute data (inclusive) = 5 steps
+    assert sample["generation_input"].shape == (*batch_dim, 3)
+    assert sample["generation_target"].shape == (*batch_dim, 5)
     # The config uses 3 periods each of which generates a sin and cos embedding
     assert sample["t0_embedding"].shape == (*batch_dim, 6)
 
@@ -184,11 +205,11 @@ def test_solar_position_decoupling(tmp_path, pvnet_config_filename):
     """Test that solar position calculations are properly decoupled from data sources."""
     config = load_yaml_configuration(pvnet_config_filename)
     config_without_solar = config.model_copy(deep=True)
-    config_without_solar.input_data.solar_position = None
+    config_without_solar.solar_position = None
 
     # Create version with explicit solar position configuration
     config_with_solar = config.model_copy(deep=True)
-    config_with_solar.input_data.solar_position = SolarPosition(
+    config_with_solar.solar_position = SolarPosition(
         time_resolution_minutes=30,
         interval_start_minutes=0,
         interval_end_minutes=180,
@@ -217,6 +238,29 @@ def test_solar_position_decoupling(tmp_path, pvnet_config_filename):
         assert key in sample_with_solar, f"Solar key {key} should be in sample"
 
 
+def test_pvnet_dataset_without_generation(tmp_path, pvnet_config_filename):
+    """Test that a dataset can be built with locations/NWP/satellite but no generation at all."""
+    config = load_yaml_configuration(pvnet_config_filename)
+    config.generation = None
+
+    config_path = tmp_path / "config_without_generation.yaml"
+    save_yaml_configuration(config, config_path)
+
+    dataset = PVNetDataset(config_path)
+
+    # With no generation data, there's nothing to be incomplete about
+    assert dataset.complete_generation
+
+    # All locations from the locations catalog are available - none to filter out
+    assert len(dataset.locations) == 317
+
+    sample = dataset[0]
+    assert "generation_input" not in sample
+    assert "generation_target" not in sample
+    assert "nwp_ukv" in sample
+    assert "satellite" in sample
+
+
 def test_pvnet_dataset_raw_sample_iteration(pvnet_config_filename):
     """Tests iterating raw samples (dict of tensors) from PVNetDataset"""
     dataset = PVNetDataset(pvnet_config_filename)
@@ -239,9 +283,14 @@ def test_pvnet_dataset_raw_sample_iteration(pvnet_config_filename):
     required_keys = [
         "nwp_ukv",
         "satellite",
-        "generation",
+        "generation_input",
+        "generation_target",
         "solar_azimuth",
         "solar_elevation",
+        "date_sin",
+        "date_cos",
+        "time_sin",
+        "time_cos",
         "location_id",
     ]
     for key in required_keys:
@@ -249,7 +298,8 @@ def test_pvnet_dataset_raw_sample_iteration(pvnet_config_filename):
 
     # Type assertions
     assert isinstance(raw_sample["satellite"], torch.Tensor)
-    assert isinstance(raw_sample["generation"], torch.Tensor)
+    assert isinstance(raw_sample["generation_input"], torch.Tensor)
+    assert isinstance(raw_sample["generation_target"], torch.Tensor)
     assert isinstance(raw_sample["solar_azimuth"], torch.Tensor)
     assert isinstance(raw_sample["solar_elevation"], torch.Tensor)
     assert isinstance(raw_sample["nwp_ukv"], torch.Tensor)
@@ -257,16 +307,25 @@ def test_pvnet_dataset_raw_sample_iteration(pvnet_config_filename):
     # Shape assertions
     assert raw_sample["satellite"].shape == (7, 1, 2, 2)
     assert raw_sample["nwp_ukv"].shape == (4, 1, 2, 2)
-    assert raw_sample["generation"].shape == (7,)
+    assert raw_sample["generation_input"].shape == (3,)
+    assert raw_sample["generation_target"].shape == (5,)
 
     # Solar position shapes - no batch dimension
-    solar_config = dataset.config.input_data.solar_position
+    solar_config = dataset.config.solar_position
     expected_time_steps = (
         solar_config.interval_end_minutes - solar_config.interval_start_minutes
     ) // solar_config.time_resolution_minutes + 1
     assert raw_sample["solar_azimuth"].shape == (expected_time_steps,)
     assert raw_sample["solar_elevation"].shape == (expected_time_steps,)
 
+    # Datetime encoding shapes - no batch dimension
+    dt_config = dataset.config.datetime_encoding
+    expected_dt_time_steps = (
+        dt_config.interval_end_minutes - dt_config.interval_start_minutes
+    ) // dt_config.time_resolution_minutes + 1
+    for key in ("date_sin", "date_cos", "time_sin", "time_cos"):
+        assert raw_sample[key].shape == (expected_dt_time_steps,)
+
     assert isinstance(raw_sample["location_id"], int | np.integer)
 
 
diff --git a/tests/datasets/pvnet/test_preprocess.py b/tests/datasets/pvnet/test_preprocess.py
index 63dd9e9c..ea89fa99 100644
--- a/tests/datasets/pvnet/test_preprocess.py
+++ b/tests/datasets/pvnet/test_preprocess.py
@@ -5,6 +5,7 @@
 from ocf_data_sampler.datasets.pvnet.preprocess import (
     apply_dropout_to_datasets,
     fill_nans_in_dataset_dicts,
+    normalise_dataset_dicts,
 )
 
 
@@ -14,34 +15,74 @@ def test_fill_nans_in_dataset_dicts(config_filename):
     configuration = load_yaml_configuration(config_filename)
 
     # Set custom satellite and nwp values, generation is left as default 0.0
-    configuration.input_data.satellite.dropout_fill_value = -1.0
-    configuration.input_data.nwp["ukv"].dropout_fill_value = -2.0
+    configuration.satellite.dropout_fill_value = -1.0
+    configuration.nwp["ukv"].dropout_fill_value = -2.0
 
     gen = np.array([1.0, np.nan, 3.0, np.nan])
     sat = np.array([1.0, np.nan, 3.0, np.nan])
     ukv = np.array([np.nan, 3.0, np.nan])
 
     datasets_dict = {
-        "generation": xr.DataArray(gen),
+        "generation_input": xr.DataArray(gen),
+        "generation_target": xr.DataArray(gen.copy()),
         "sat": xr.DataArray(sat),
         "nwp": {"ukv": xr.DataArray(ukv)},
     }
 
     datasets_dict = fill_nans_in_dataset_dicts(datasets_dict, config=configuration)
 
-    assert np.array_equal(datasets_dict["generation"].values, np.array([1.0, 0.0, 3.0, 0.0]))
+    expected_gen = np.array([1.0, 0.0, 3.0, 0.0])
+    assert np.array_equal(datasets_dict["generation_input"].values, expected_gen)
+    assert np.array_equal(datasets_dict["generation_target"].values, expected_gen)
     assert np.array_equal(datasets_dict["sat"].values, np.array([1.0, -1.0, 3.0, -1.0]))
     assert np.array_equal(datasets_dict["nwp"]["ukv"].values, np.array([-2.0, 3.0, -2.0]))
 
 
+def test_normalise_dataset_dicts_generation():
+    """Generation is normalised to a capacity factor (generation_mw / capacity_mwp)."""
+    generation_mw = np.array([[50.0, 0.0], [100.0, 20.0]])
+    capacity_mwp = np.array([[100.0, 0.0], [100.0, 40.0]])
+    generation = xr.DataArray(
+        np.stack([generation_mw, capacity_mwp], axis=-1),
+        coords={
+            "time_utc": ["2023-01-01T00:00", "2023-01-01T00:30"],
+            "location_id": [1, 2],
+            "gen_param": ["generation_mw", "capacity_mwp"],
+        },
+        dims=("time_utc", "location_id", "gen_param"),
+    )
+
+    datasets_dict = {
+        "generation_input": generation,
+        "generation_target": generation.copy(deep=True),
+    }
+
+    datasets_dict = normalise_dataset_dicts(datasets_dict, {}, {}, {}, {})
+
+    for key in ("generation_input", "generation_target"):
+        result = datasets_dict[key]
+
+        # generation_mw is rescaled element-wise by capacity_mwp
+        normalised = result.sel(gen_param="generation_mw").values
+        assert normalised[0, 0] == 0.5
+        assert normalised[1, 0] == 1.0
+        assert normalised[1, 1] == 0.5
+
+        # Zero capacity normalises to NaN rather than raw MW or a division error
+        assert np.isnan(normalised[0, 1])
+
+        # capacity_mwp itself is left unchanged
+        assert np.array_equal(result.sel(gen_param="capacity_mwp").values, capacity_mwp)
+
+
 def test_apply_dropout_to_datasets(pvnet_config_filename):
     config = load_yaml_configuration(pvnet_config_filename)
 
-    # Set dropout
-    config.input_data.generation.dropout_timedeltas_minutes = [-30]
-    config.input_data.generation.dropout_fraction = 1.0
-    config.input_data.satellite.dropout_timedeltas_minutes = []
-    config.input_data.satellite.dropout_fraction = 0
+    # Set dropout on the input window only - generation.target has no dropout config
+    config.generation.input.dropout_timedeltas_minutes = [-30]
+    config.generation.input.dropout_fraction = 1.0
+    config.satellite.dropout_timedeltas_minutes = []
+    config.satellite.dropout_fraction = 0
 
     t0 = np.datetime64("2023-01-01 12:00")
     times = np.array(
@@ -53,10 +94,16 @@ def test_apply_dropout_to_datasets(pvnet_config_filename):
         ],
         dtype="datetime64[m]",
     )
+    generation_mw = np.arange(4 * 2, dtype=float).reshape(4, 2)
+    capacity_mwp = np.ones((4, 2))
     generation = xr.DataArray(
-        np.arange(4 * 2, dtype=float).reshape(4, 2),
-        coords={"time_utc": times, "location_id": [1, 2]},
-        dims=("time_utc", "location_id"),
+        np.stack([generation_mw, capacity_mwp], axis=-1),
+        coords={
+            "time_utc": times,
+            "location_id": [1, 2],
+            "gen_param": ["generation_mw", "capacity_mwp"],
+        },
+        dims=("time_utc", "location_id", "gen_param"),
     )
     sat = xr.DataArray(
         np.arange(4, dtype=float),
@@ -64,16 +111,21 @@ def test_apply_dropout_to_datasets(pvnet_config_filename):
         dims=("time_utc",),
     )
 
-    datasets_dict = {"generation": generation, "sat": sat}
+    datasets_dict = {"generation_input": generation, "sat": sat}
 
     apply_dropout_to_datasets(datasets_dict, t0, config)
 
-    ds_gen = datasets_dict["generation"]
+    ds_gen = datasets_dict["generation_input"].sel(gen_param="generation_mw")
+    ds_cap = datasets_dict["generation_input"].sel(gen_param="capacity_mwp")
 
-    # Generation dropout with a -30 minute history should blank only the t0 timestep.
+    # Generation dropout with a -30 minute cutoff should blank everything from t0 onwards,
+    # including timesteps beyond t0.
     assert not np.any(np.isnan(ds_gen.sel(time_utc=slice(None, "2023-01-01T11:30"))))
-    assert np.all(np.isnan(ds_gen.sel(time_utc=t0)))
-    assert not np.any(np.isnan(ds_gen.sel(time_utc=slice("2023-01-01T12:30", None))))
+    assert np.all(np.isnan(ds_gen.sel(time_utc=slice("2023-01-01T12:00", None))))
+
+    # capacity_mwp is dropped out along with generation_mw, using the same cutoff.
+    assert not np.any(np.isnan(ds_cap.sel(time_utc=slice(None, "2023-01-01T11:30"))))
+    assert np.all(np.isnan(ds_cap.sel(time_utc=slice("2023-01-01T12:00", None))))
 
     # Satellite dropout is disabled, so the helper should leave it untouched.
     xr.testing.assert_equal(datasets_dict["sat"], sat)
diff --git a/tests/datasets/pvnet/test_sample.py b/tests/datasets/pvnet/test_sample.py
index d876f515..33d5dd93 100644
--- a/tests/datasets/pvnet/test_sample.py
+++ b/tests/datasets/pvnet/test_sample.py
@@ -19,37 +19,38 @@ def test_make_sun_position_numpy_sample():
 
 
 def test_convert_generation_to_numpy_sample(generation_zarr_path):
+    """convert_to_numpy_sample just extracts generation_mw/capacity_mwp as-is - normalising
+    generation_mw to a capacity factor is normalise_dataset_dicts's job (see test_preprocess.py).
+    """
     da = open_generation(generation_zarr_path).isel(time_utc=slice(0, 10)).sel(location_id=1)
-    t0_idx = 0
-    numpy_sample = convert_to_numpy_sample({"generation": da}, t0_idx=t0_idx)
+    numpy_sample = convert_to_numpy_sample({"generation_input": da, "generation_target": da})
 
-    # Assert structure
-    assert isinstance(numpy_sample, dict)
-    assert "generation" in numpy_sample
-    assert "capacity_mwp" in numpy_sample
-    assert "generation_time_utc" in numpy_sample
+    generation_mw = da.sel(gen_param="generation_mw").values
+    capacity = da.sel(gen_param="capacity_mwp").values
+
+    for key in ("generation_input", "generation_target"):
+        # Assert structure
+        assert isinstance(numpy_sample, dict)
+        assert key in numpy_sample
+        assert f"{key}_capacity_mwp" in numpy_sample
+        assert f"{key}_time_utc" in numpy_sample
 
-    # Assert content and capacity values
-    assert np.array_equal(numpy_sample["generation"], da.sel(gen_param="generation_mw").values)
-    assert isinstance(numpy_sample["generation_time_utc"], np.ndarray)
-    assert numpy_sample["generation_time_utc"].dtype == float
-    assert numpy_sample["capacity_mwp"] == da.sel(gen_param="capacity_mwp").isel(time_utc=0).values
+        # Assert content is passed through unchanged
+        assert np.array_equal(numpy_sample[key], generation_mw)
+        assert np.array_equal(numpy_sample[f"{key}_capacity_mwp"], capacity)
+        assert isinstance(numpy_sample[f"{key}_time_utc"], np.ndarray)
+        assert numpy_sample[f"{key}_time_utc"].dtype == float
 
 
 def test_convert_nwp_to_numpy_sample(ds_nwp_ukv_time_sliced):
-    t0_idx = 0
-    numpy_sample = convert_to_numpy_sample(
-        {"nwp": {"ukv": ds_nwp_ukv_time_sliced}},
-        t0_idx=t0_idx,
-    )
+    numpy_sample = convert_to_numpy_sample({"nwp": {"ukv": ds_nwp_ukv_time_sliced}})
 
     assert isinstance(numpy_sample, dict)
     assert (numpy_sample["nwp_ukv"] == ds_nwp_ukv_time_sliced.values).all()
 
 
 def test_convert_satellite_to_numpy_sample(da_sat_like):
-    t0_idx = 0
-    numpy_sample = convert_to_numpy_sample({"sat": da_sat_like}, t0_idx=t0_idx)
+    numpy_sample = convert_to_numpy_sample({"sat": da_sat_like})
 
     assert isinstance(numpy_sample, dict)
     assert (numpy_sample["satellite"] == da_sat_like.values).all()
diff --git a/tests/fixtures/configs/pvnet_test_config.yaml b/tests/fixtures/configs/pvnet_test_config.yaml
index 38664772..cf868500 100644
--- a/tests/fixtures/configs/pvnet_test_config.yaml
+++ b/tests/fixtures/configs/pvnet_test_config.yaml
@@ -1,58 +1,65 @@
-general:
-  description: Test config for PVNet
-  name: pvnet_test
+sampling_grid:
+  locations_zarr_path: set_in_temp_file
+  t0_resolution_minutes: 30
 
-input_data:
-
-  generation:
-    zarr_path: set_in_temp_file
+generation:
+  zarr_path: set_in_temp_file
+  time_resolution_minutes: 30
+  input:
     interval_start_minutes: -60
-    interval_end_minutes: 120
-    time_resolution_minutes: 30
+    interval_end_minutes: 0
     dropout_timedeltas_minutes: []
     dropout_fraction: 0
+  target:
+    interval_start_minutes: 0
+    interval_end_minutes: 120
 
-  nwp:
-    ukv:
-      provider: ukv
-      zarr_path: set_in_temp_file
-      interval_start_minutes: -60
-      interval_end_minutes: 120
-      time_resolution_minutes: 60
-      channels:
-        - t # 2-metre temperature
-      image_size_pixels_height: 2
-      image_size_pixels_width: 2
-      dropout_timedeltas_minutes: [-180]
-      dropout_fraction: 1.0
-      max_staleness_minutes: null
-      normalisation_constants:
-        t:
-          mean: 283.64913206
-          std: 4.38818501
-          clip_min: 270
-          clip_max: 310
-
-  satellite:
+nwp:
+  ukv:
+    provider: ukv
     zarr_path: set_in_temp_file
-    interval_start_minutes: -30
-    interval_end_minutes: 0
-    time_resolution_minutes: 5
+    interval_start_minutes: -60
+    interval_end_minutes: 120
+    time_resolution_minutes: 60
     channels:
-      - IR_016
+      - t # 2-metre temperature
     image_size_pixels_height: 2
     image_size_pixels_width: 2
-    dropout_timedeltas_minutes: []
-    dropout_fraction: 0
+    dropout_timedeltas_minutes: [-180]
+    dropout_fraction: 1.0
+    max_staleness_minutes: null
     normalisation_constants:
-      IR_016:
-        mean: 0.17594202
-        std: 0.21462157
+      t:
+        mean: 283.64913206
+        std: 4.38818501
+        clip_min: 270
+        clip_max: 310
+
+satellite:
+  zarr_path: set_in_temp_file
+  interval_start_minutes: -30
+  interval_end_minutes: 0
+  time_resolution_minutes: 5
+  channels:
+    - IR_016
+  image_size_pixels_height: 2
+  image_size_pixels_width: 2
+  dropout_timedeltas_minutes: []
+  dropout_fraction: 0
+  normalisation_constants:
+    IR_016:
+      mean: 0.17594202
+      std: 0.21462157
+
+solar_position:
+  interval_start_minutes: -15
+  interval_end_minutes: 15
+  time_resolution_minutes: 5
 
-  solar_position:
-    interval_start_minutes: -15
-    interval_end_minutes: 15
-    time_resolution_minutes: 5
+datetime_encoding:
+  interval_start_minutes: -15
+  interval_end_minutes: 15
+  time_resolution_minutes: 5
 
-  t0_embedding:
-    embeddings: [["1h", "cyclic"], ["24h", "cyclic"], ["1y", "cyclic"]]
+t0_embedding:
+  embeddings: [["1h", "cyclic"], ["24h", "cyclic"], ["1y", "cyclic"]]
diff --git a/tests/load/test_load_generation.py b/tests/load/test_load_generation.py
index c425a2f1..56c5dc04 100644
--- a/tests/load/test_load_generation.py
+++ b/tests/load/test_load_generation.py
@@ -14,7 +14,6 @@ def test_open_generation(generation_zarr_path):
 
     assert isinstance(da, xr.DataArray)
     assert da.dims == ("time_utc", "location_id", "gen_param")
-    assert {"longitude", "latitude"}.issubset(da.coords)
     assert da.shape == (49, 318, 2)
     assert len(np.unique(da.coords["location_id"])) == da.shape[1]
 
diff --git a/tests/load/test_load_locations.py b/tests/load/test_load_locations.py
new file mode 100644
index 00000000..e4bed3cb
--- /dev/null
+++ b/tests/load/test_load_locations.py
@@ -0,0 +1,56 @@
+from pathlib import Path
+
+import numpy as np
+import pytest
+import xarray as xr
+
+from ocf_data_sampler.load.locations import open_locations
+
+
+def test_open_locations(locations_zarr_path):
+    """Test the locations data loader with valid data."""
+    ds = open_locations(locations_zarr_path)
+
+    assert isinstance(ds, xr.Dataset)
+    assert set(ds.data_vars) == {"longitude", "latitude"}
+    assert ds["longitude"].dims == ("location_id",)
+    assert ds["latitude"].dims == ("location_id",)
+    assert len(np.unique(ds.coords["location_id"])) == ds.sizes["location_id"]
+
+
+def test_open_locations_missing_data_var(tmp_path: Path):
+    """Test that open_locations raises a ValueError when a required data variable is missing."""
+    zarr_path = tmp_path / "bad_locations.zarr"
+
+    bad_ds = xr.Dataset(
+        data_vars={
+            "longitude": (("location_id",), [0.0, 1.0]),
+        },
+        coords={
+            "location_id": [1, 2],
+        },
+    )
+    bad_ds.to_zarr(zarr_path)
+
+    with pytest.raises(ValueError, match="Locations data should have variables"):
+        open_locations(zarr_path=str(zarr_path))
+
+
+def test_open_locations_bad_dtype(tmp_path: Path):
+    """Test that open_locations raises a TypeError on incorrect data dtypes."""
+    zarr_path = tmp_path / "bad_locations.zarr"
+
+    # Create dataset where longitude is integer
+    bad_ds = xr.Dataset(
+        data_vars={
+            "longitude": (("location_id",), [0, 1]),
+            "latitude": (("location_id",), [0.0, 1.0]),
+        },
+        coords={
+            "location_id": [1, 2],
+        },
+    )
+    bad_ds.to_zarr(zarr_path)
+
+    with pytest.raises(TypeError, match="longitude in locations data should be floating"):
+        open_locations(zarr_path=str(zarr_path))
diff --git a/tests/select/test_dropout.py b/tests/select/test_dropout.py
index 4652f31d..5220d226 100644
--- a/tests/select/test_dropout.py
+++ b/tests/select/test_dropout.py
@@ -3,17 +3,17 @@
 import xarray as xr
 
 from ocf_data_sampler.common.time_utils import minutes
-from ocf_data_sampler.select.dropout import apply_history_dropout
+from ocf_data_sampler.select.dropout import apply_dropout
 
 
-def test_apply_history_dropout_multiple_timedeltas(da_sample):
+def test_apply_dropout_multiple_timedeltas(da_sample):
 
     # Dropout edits the input in-place, so make a copy to avoid affecting other tests
     da_sample = da_sample.copy(deep=True)
 
     t0 = da_sample["time_utc"].values[-1]
 
-    da_sample_dropout = apply_history_dropout(
+    da_sample_dropout = apply_dropout(
         da_sample,
         t0,
         dropout_timedeltas=minutes([-30, -45]),
@@ -30,14 +30,14 @@ def test_apply_history_dropout_multiple_timedeltas(da_sample):
     )
 
 
-def test_apply_history_dropout_none(da_sample):
+def test_apply_dropout_none(da_sample):
 
     # Dropout edits the input in-place, so make a copy to avoid affecting other tests
     da_sample = da_sample.copy(deep=True)
 
     t0 = da_sample["time_utc"].values[-1]
 
-    da_sample_dropout = apply_history_dropout(
+    da_sample_dropout = apply_dropout(
         da_sample,
         t0,
         dropout_timedeltas=[minutes(-30)],
@@ -45,7 +45,7 @@ def test_apply_history_dropout_none(da_sample):
     )
     xr.testing.assert_equal(da_sample_dropout, da_sample)
 
-    da_sample_dropout = apply_history_dropout(
+    da_sample_dropout = apply_dropout(
         da_sample,
         t0,
         dropout_timedeltas=[],
@@ -54,14 +54,14 @@ def test_apply_history_dropout_none(da_sample):
     xr.testing.assert_equal(da_sample_dropout, da_sample)
 
 
-def test_apply_history_dropout_list(da_sample):
+def test_apply_dropout_list(da_sample):
 
     # Dropout edits the input in-place, so make a copy to avoid affecting other tests
     da_sample = da_sample.copy(deep=True)
 
     t0 = da_sample["time_utc"].values[-1]
 
-    da_sample_dropout = apply_history_dropout(
+    da_sample_dropout = apply_dropout(
         da_sample,
         t0,
         dropout_timedeltas=minutes([-30, -45]),
@@ -79,7 +79,7 @@ def test_apply_history_dropout_list(da_sample):
 
 
 @pytest.mark.parametrize("t0_str", ["12:30", "13:00", "13:30"])
-def test_apply_history_dropout(da_sample, t0_str):
+def test_apply_dropout(da_sample, t0_str):
 
     # Dropout edits the input in-place, so make a copy to avoid affecting other tests
     da_sample = da_sample.copy(deep=True)
@@ -87,12 +87,13 @@ def test_apply_history_dropout(da_sample, t0_str):
     t0_time = pd.Timestamp(f"2024-01-01 {t0_str}")
     dropout_time = t0_time + minutes(-30)
 
-    da_dropout = apply_history_dropout(
+    da_dropout = apply_dropout(
         da_sample,
         t0_time,
         dropout_timedeltas=[minutes(-30)],
         dropout_frac=1.0,
     )
 
+    # Everything after the dropout cut-off is masked, including any data beyond t0.
     assert da_dropout.sel(time_utc=slice(None, dropout_time)).notnull().all()
-    assert da_dropout.sel(time_utc=slice(dropout_time + minutes(5), t0_time)).isnull().all()
+    assert da_dropout.sel(time_utc=slice(dropout_time + minutes(5), None)).isnull().all()

From ca98c8f0da96a025608a053d849f85d01bb121a8 Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Thu, 30 Jul 2026 11:20:12 +0000
Subject: [PATCH 2/9] clean up

---
 src/ocf_data_sampler/config/model.py          | 42 +++++--------
 .../datasets/pvnet/dataset.py                 |  4 +-
 .../datasets/pvnet/loading.py                 |  9 +--
 .../datasets/pvnet/preprocess.py              | 61 +++++++++++--------
 tests/datasets/pvnet/test_preprocess.py       | 55 ++++++++++++++++-
 5 files changed, 107 insertions(+), 64 deletions(-)

diff --git a/src/ocf_data_sampler/config/model.py b/src/ocf_data_sampler/config/model.py
index 39611182..5ae27cc7 100644
--- a/src/ocf_data_sampler/config/model.py
+++ b/src/ocf_data_sampler/config/model.py
@@ -310,11 +310,10 @@ def items(self) -> Iterator[tuple[str, NWP]]:
 
 
 class GenerationWindow(Base):
-    """Mixin class, to add interval start and end minutes for a generation window.
+    """Interval bounds for a generation window.
 
-    Unlike `TimeWindowMixin`, the temporal resolution is not included here - it belongs to the
-    shared generation data source (`Generation.time_resolution_minutes`), not to an individual
-    window over it.
+    Unlike `TimeWindowMixin`, no temporal resolution here - it belongs to the shared
+    generation data source (`Generation.time_resolution_minutes`), not to an individual window.
     """
 
     interval_start_minutes: int = Field(
@@ -340,27 +339,19 @@ def validate_interval_order(self) -> "GenerationWindow":
 
 
 class GenerationInputWindow(GenerationWindow, DropoutMixin):
-    """Generation input window configuration model, used for `Generation.input`.
+    """A generation window with dropout configuration.
 
-    Extends `GenerationWindow` with dropout configuration, since only the input window (not the
-    prediction target) should ever be randomly masked out.
+    Dropout is configurable here since only the input window (not the prediction target)
+    should ever be randomly masked out.
     """
 
 
 class GenerationTargetWindow(GenerationWindow, FillValueMixin):
-    """Generation target window configuration model, used for `Generation.target`."""
+    """Generation target window configuration model."""
 
 
 class Generation(Base):
-    """Generation configuration model.
-
-    Bundles the shared generation data source (`zarr_path`, `time_resolution_minutes`) with its
-    `input` and `target` windows - two independently configurable time windows over the same
-    underlying data. `time_resolution_minutes` describes generation's own native data cadence
-    (used for gap detection and windowed slicing of generation's own data) - it is independent
-    of `SamplingGrid.t0_resolution_minutes`, which is the cadence t0 candidates are enumerated
-    at and may legitimately differ (e.g. generation stored every 5 minutes, sampled every 30).
-    """
+    """Generation configuration model."""
 
     zarr_path: str = Field(
         ...,
@@ -401,25 +392,20 @@ def validate_windows(self) -> "Generation":
 
 
 class SamplingGrid(Base):
-    """Configuration for the (location, time) grid that t0 times are sampled from.
-
-    `locations_zarr_path` points to the locations metadata (location IDs and their
-    coordinates) - see `ocf_data_sampler.load.locations.open_locations`.
-    `t0_resolution_minutes` is the cadence t0 candidates are enumerated at, needed to compute
-    valid t0 times regardless of which other input sources are configured - it is not any one
-    source's own native data resolution (see `Generation.time_resolution_minutes` for that).
-    """
+    """Configuration for the (location, time) grid that t0 times are sampled from."""
 
     locations_zarr_path: str = Field(
         ...,
-        description="Absolute or relative zarr filepath to the locations metadata. Prefix with "
-        "a protocol like s3:// to read from alternative filesystems.",
+        description="Absolute or relative zarr filepath to the locations metadata (location IDs "
+        "and their coordinates) - see `ocf_data_sampler.load.locations.open_locations`. Prefix "
+        "with a protocol like s3:// to read from alternative filesystems.",
     )
 
     t0_resolution_minutes: int = Field(
         ...,
         gt=0,
-        description="The resolution of the t0 sampling grid, in minutes.",
+        description="The cadence t0 candidates are enumerated at, needed to compute valid t0 "
+        "times regardless of which other input sources are configured.",
     )
 
 
diff --git a/src/ocf_data_sampler/datasets/pvnet/dataset.py b/src/ocf_data_sampler/datasets/pvnet/dataset.py
index 27895212..e686927f 100644
--- a/src/ocf_data_sampler/datasets/pvnet/dataset.py
+++ b/src/ocf_data_sampler/datasets/pvnet/dataset.py
@@ -265,8 +265,8 @@ def __init__(
             if len(missing) > 0:
                 raise ValueError(f"Generation data is missing for location IDs: {missing}")
 
-            # Restrict to the catalog's locations - generation may have extra ids that aren't
-            # real samplable points (e.g. a summation-model placeholder).
+            # Slice the generation data to only include the specified locations. This allows us to
+            # quality check the generation data for nans and find valid t0 times for each location.
             datasets_dict["generation"] = datasets_dict["generation"].sel(location_id=location_ids)
 
         # Check if generation data has nans. If generation isn't configured at all, there's no
diff --git a/src/ocf_data_sampler/datasets/pvnet/loading.py b/src/ocf_data_sampler/datasets/pvnet/loading.py
index ecd63f84..63e4e5f8 100644
--- a/src/ocf_data_sampler/datasets/pvnet/loading.py
+++ b/src/ocf_data_sampler/datasets/pvnet/loading.py
@@ -10,18 +10,15 @@
 def get_dataset_dict(config: PVNetDataConfig) -> SourceDict[xr.DataArray]:
     """Construct dictionary of all of the per-sample input data sources.
 
-    Locations metadata is not included here - unlike generation/nwp/satellite it isn't a
-    per-sample source, so it's loaded separately by the caller - see
-    `ocf_data_sampler.load.locations.open_locations`.
+    Locations metadata is deliberately excluded - it isn't a per-sample source, so the caller
+    loads it separately.
 
     Args:
         config: PVNetDataConfig configuration object
     """
     datasets_dict = {}
 
-    # Load generation data unless not configured. Any locations generation has that aren't in
-    # the locations catalog (e.g. an id used only for summation models) get filtered out later,
-    # once the catalog is available - see AbstractPVNetDataset.__init__.
+    # Load generation data if in config
     if config.generation is not None:
         datasets_dict["generation"] = open_generation(zarr_path=config.generation.zarr_path)
 
diff --git a/src/ocf_data_sampler/datasets/pvnet/preprocess.py b/src/ocf_data_sampler/datasets/pvnet/preprocess.py
index 2759b7bc..ca8462e4 100644
--- a/src/ocf_data_sampler/datasets/pvnet/preprocess.py
+++ b/src/ocf_data_sampler/datasets/pvnet/preprocess.py
@@ -92,12 +92,6 @@ def normalise_dataset_dicts(
 ) -> SourceDict:
     """Normalise the NWP, satellite, and generation data in-place.
 
-    NWP and satellite are normalised using the per-channel mean/std/clip constants from config.
-    Generation is normalised differently: `generation_mw` is rescaled to a capacity factor by
-    dividing by `capacity_mwp`, which is time-varying, per-location data rather than a config
-    constant - so it can't use the same clip/mean/std path. `capacity_mwp` itself is left
-    unchanged, since it's exposed raw in the output sample.
-
     Args:
         dataset_dict: Dictionary of xarray datasets
         mean_dict: Means, as constructed by `config_normalization_values_to_dicts`
@@ -127,30 +121,45 @@ def normalise_dataset_dicts(
         )
 
     for key in ("generation_input", "generation_target"):
-        if key not in dataset_dict:
-            continue
-
-        da = dataset_dict[key]
-        gen_idx = list(da["gen_param"].values).index("generation_mw")
-        cap_idx = list(da["gen_param"].values).index("capacity_mwp")
-
-        generation_values = da.isel(gen_param=gen_idx).values
-        capacity_values = da.isel(gen_param=cap_idx).values
-
-        # capacity_mwp is time-varying (per timestep, per location) - normalise element-wise
-        # rather than by a single scalar. Where capacity is 0 the ratio is undefined, so we
-        # emit NaN rather than silently switching units (raw MW) or dividing by zero - dropout
-        # fill (later in the pipeline) replaces it with a fixed value in normalised units.
-        da.data[..., gen_idx] = np.divide(
-            generation_values,
-            capacity_values,
-            out=np.full_like(generation_values, np.nan, dtype=float),
-            where=capacity_values != 0,
-        )
+        if key in dataset_dict:
+            dataset_dict[key] = normalise_generation_by_capacity(dataset_dict[key])
 
     return dataset_dict
 
 
+def normalise_generation_by_capacity(da: TArray) -> TArray:
+    """Rescale `generation_mw` to a capacity factor, leaving `capacity_mwp` unchanged.
+
+    Zero capacity means no plant, so the capacity factor is taken as 0 rather than the undefined
+    0/0. Emitting NaN instead would route it through the dropout fill, which signals missing data
+    rather than a known-zero output.
+
+    Args:
+        da: Generation DataArray-like with a `gen_param` dimension
+    """
+    gen_params = list(da["gen_param"].values)
+    gen_idx = gen_params.index("generation_mw")
+    cap_idx = gen_params.index("capacity_mwp")
+
+    generation_values = da.isel(gen_param=gen_idx).values
+    capacity_values = da.isel(gen_param=cap_idx).values
+
+    normalised = np.divide(
+        generation_values,
+        capacity_values,
+        out=np.zeros_like(generation_values, dtype=float),
+        where=capacity_values != 0,
+    )
+
+    new_data = da.data.copy()
+    index = [slice(None)] * new_data.ndim
+    index[da.dims.index("gen_param")] = gen_idx
+    new_data[tuple(index)] = normalised
+    da.data = new_data
+
+    return da
+
+
 def diff_nwp_data(dataset_dict: SourceDict, config: PVNetDataConfig) -> SourceDict:
     """Take the in-place diff of some channels of the NWP data.
 
diff --git a/tests/datasets/pvnet/test_preprocess.py b/tests/datasets/pvnet/test_preprocess.py
index ea89fa99..f6040c46 100644
--- a/tests/datasets/pvnet/test_preprocess.py
+++ b/tests/datasets/pvnet/test_preprocess.py
@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 import xarray as xr
 
 from ocf_data_sampler.config import load_yaml_configuration
@@ -68,13 +69,63 @@ def test_normalise_dataset_dicts_generation():
         assert normalised[1, 0] == 1.0
         assert normalised[1, 1] == 0.5
 
-        # Zero capacity normalises to NaN rather than raw MW or a division error
-        assert np.isnan(normalised[0, 1])
+        # Zero capacity normalises to 0 rather than raw MW or a division error
+        assert normalised[0, 1] == 0.0
 
         # capacity_mwp itself is left unchanged
         assert np.array_equal(result.sel(gen_param="capacity_mwp").values, capacity_mwp)
 
 
+def _generation_source(generation_mw: np.ndarray, capacity_mwp: np.ndarray) -> xr.DataArray:
+    """Build a generation DataArray with `n` timesteps for a single location."""
+    return xr.DataArray(
+        np.stack([generation_mw[:, None], capacity_mwp[:, None]], axis=-1),
+        coords={
+            "time_utc": pd.date_range("2023-01-01", periods=len(generation_mw), freq="30min"),
+            "location_id": [1],
+            "gen_param": ["generation_mw", "capacity_mwp"],
+        },
+        dims=("time_utc", "location_id", "gen_param"),
+    )
+
+
+def test_normalise_does_not_mutate_shared_source():
+    """Normalising windows must not write through to the array they were sliced from.
+
+    Time slices are views onto the eagerly-loaded source, so an in-place write would corrupt it
+    for every later sample and double-normalise timesteps shared by the two windows.
+    """
+    src = _generation_source(np.array([100.0, 150.0, 200.0]), np.full(3, 100.0))
+    before = src.values.copy()
+
+    # Windows deliberately overlap on the middle timestep
+    datasets_dict = {
+        "generation_input": src.isel(time_utc=slice(0, 2)),
+        "generation_target": src.isel(time_utc=slice(1, 3)),
+    }
+    datasets_dict = normalise_dataset_dicts(datasets_dict, {}, {}, {}, {})
+
+    assert np.array_equal(src.values, before), "source array was mutated by normalisation"
+
+    def generation_of(key: str) -> np.ndarray:
+        return datasets_dict[key].sel(gen_param="generation_mw").values.ravel()
+
+    assert np.array_equal(generation_of("generation_input"), [1.0, 1.5])
+    # The shared timestep is normalised once, not once per window
+    assert np.array_equal(generation_of("generation_target"), [1.5, 2.0])
+
+
+def test_normalise_generation_is_dimension_order_agnostic():
+    """`gen_param` is indexed by name, so it need not be the last dimension."""
+    src = _generation_source(np.array([100.0, 150.0, 200.0]), np.full(3, 100.0))
+    transposed = src.transpose("gen_param", "time_utc", "location_id")
+
+    result = normalise_dataset_dicts({"generation_input": transposed}, {}, {}, {}, {})
+
+    normalised = result["generation_input"].sel(gen_param="generation_mw").values.ravel()
+    assert np.array_equal(normalised, [1.0, 1.5, 2.0])
+
+
 def test_apply_dropout_to_datasets(pvnet_config_filename):
     config = load_yaml_configuration(pvnet_config_filename)
 

From 34973f6967ddee786313847182f22c2131e4e3bd Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Thu, 30 Jul 2026 11:52:11 +0000
Subject: [PATCH 3/9] clean up more

---
 .../datasets/pvnet/dataset.py                 | 26 +++---
 .../datasets/pvnet/preprocess.py              |  5 +-
 src/ocf_data_sampler/datasets/pvnet/sample.py |  4 +-
 .../datasets/pvnet/slicing.py                 | 19 ++--
 tests/conftest.py                             | 90 +++++++++----------
 tests/datasets/pvnet/test_preprocess.py       | 69 +++-----------
 tests/load/test_load_locations.py             | 20 -----
 7 files changed, 82 insertions(+), 151 deletions(-)

diff --git a/src/ocf_data_sampler/datasets/pvnet/dataset.py b/src/ocf_data_sampler/datasets/pvnet/dataset.py
index e686927f..8fdedf9a 100644
--- a/src/ocf_data_sampler/datasets/pvnet/dataset.py
+++ b/src/ocf_data_sampler/datasets/pvnet/dataset.py
@@ -289,7 +289,9 @@ def __init__(
                 "Generation data has nans so t0s are handled separately for each location_id.",
             )
             # If non-identical times per location, find valid t0s per location id
-            valid_t0_and_location_ids = self.find_valid_t0_and_location_ids(datasets_dict, config)
+            valid_t0_and_location_ids = self.find_valid_t0_and_location_ids(
+                datasets_dict, locations, config,
+            )
 
             # Filter t0 times to given range
             if time_periods is not None:
@@ -355,6 +357,7 @@ def find_valid_t0_times(
     @staticmethod
     def find_valid_t0_and_location_ids(
         datasets_dict: SourceDict,
+        locations: list[Location],
         config: PVNetDataConfig,
     ) -> pd.DataFrame:
         """Find the t0 times where all of the requested input data is available for each location.
@@ -365,6 +368,7 @@ def find_valid_t0_and_location_ids(
 
         Args:
             datasets_dict: A dictionary of input datasets
+            locations: The locations to find valid t0 times for
             config: PVNetDataConfig file
         """
         # Get valid time period for nwp and satellite
@@ -372,23 +376,23 @@ def find_valid_t0_and_location_ids(
         valid_time_periods = find_valid_time_periods(datasets_without_generation, config)
 
         # Loop over each location in system id and obtain valid periods
-        generations = datasets_dict["generation"]
-        location_ids = generations.location_id.values
-        generation_config = config.generation
         generation_windows = [
-            w for w in (generation_config.input, generation_config.target) if w is not None
+            w for w in (config.generation.input, config.generation.target) if w is not None
         ]
         valid_t0_and_location_ids = []
-        for location_id in location_ids:
-            generation = generations.sel(location_id=location_id)
-            # Drop NaN values
-            generation = generation.dropna(dim="time_utc")
+        for location in locations:
+            # Drop NaN values for location
+            generation = (
+                datasets_dict["generation"]
+                .sel(location_id=location.id)
+                .dropna(dim="time_utc")
+            )
 
             # Obtain valid time periods for this location, for each configured window
             time_periods_per_window = [
                 find_contiguous_t0_periods(
                     generation["time_utc"].values,
-                    time_resolution=minutes(generation_config.time_resolution_minutes),
+                    time_resolution=minutes(config.generation.time_resolution_minutes),
                     interval_start=minutes(window_config.interval_start_minutes),
                     interval_end=minutes(window_config.interval_end_minutes),
                 )
@@ -405,7 +409,7 @@ def find_valid_t0_and_location_ids(
             )
 
             valid_t0_per_location = pd.DataFrame(index=valid_t0_times_per_location)
-            valid_t0_per_location["location_id"] = location_id
+            valid_t0_per_location["location_id"] = location.id
             valid_t0_and_location_ids.append(valid_t0_per_location)
 
         valid_t0_and_location_ids = pd.concat(valid_t0_and_location_ids)
diff --git a/src/ocf_data_sampler/datasets/pvnet/preprocess.py b/src/ocf_data_sampler/datasets/pvnet/preprocess.py
index ca8462e4..ee61b54e 100644
--- a/src/ocf_data_sampler/datasets/pvnet/preprocess.py
+++ b/src/ocf_data_sampler/datasets/pvnet/preprocess.py
@@ -199,11 +199,10 @@ def apply_dropout_to_datasets(
             dropout_frac=config.satellite.dropout_fraction,
         )
 
+    # generation_target is never dropped out since it's the prediction target
     if "generation_input" in datasets_dict:
 
-        # capacity_mwp is dropped out along with generation_mw - if the input feed was stale for
-        # this timestep, we didn't know the capacity at that point either.
-        # generation_target is never dropped out - it's the prediction target, not an input.
+        # Note: capacity_mwp is dropped out along with generation_mw
         apply_dropout(
             datasets_dict["generation_input"],
             t0,
diff --git a/src/ocf_data_sampler/datasets/pvnet/sample.py b/src/ocf_data_sampler/datasets/pvnet/sample.py
index 524ddf5f..3297fcf1 100644
--- a/src/ocf_data_sampler/datasets/pvnet/sample.py
+++ b/src/ocf_data_sampler/datasets/pvnet/sample.py
@@ -34,8 +34,8 @@ def convert_to_numpy_sample(
 
         da = datasets_dict[key]
 
-        # generation_mw has already been normalised in-place to a capacity factor so should be in
-        # range [0, 1]. capacity_mwp is still in MW
+        # generation_mw is expected to already been normalised by capacity so should be in the
+        # range [0, 1]. capacity_mwp is still expected to be in MW
         gen_idx = list(da["gen_param"].values).index("generation_mw")
         cap_idx = list(da["gen_param"].values).index("capacity_mwp")
 
diff --git a/src/ocf_data_sampler/datasets/pvnet/slicing.py b/src/ocf_data_sampler/datasets/pvnet/slicing.py
index fd170560..e2a8497d 100644
--- a/src/ocf_data_sampler/datasets/pvnet/slicing.py
+++ b/src/ocf_data_sampler/datasets/pvnet/slicing.py
@@ -95,13 +95,12 @@ def reduce_spatial_extent_of_datasets(
 
 
     if "sat" in datasets_dict:
-        sat_config = config.satellite
 
         sliced_datasets_dict["sat"] = select_spatial_slice_pixels_multiple(
             datasets_dict["sat"],
             locations,
-            height_pixels=sat_config.image_size_pixels_height,
-            width_pixels=sat_config.image_size_pixels_width,
+            height_pixels=config.satellite.image_size_pixels_height,
+            width_pixels=config.satellite.image_size_pixels_width,
         )
 
     if "generation" in datasets_dict:
@@ -153,21 +152,19 @@ def slice_datasets_by_time(
             )
 
     if "sat" in datasets_dict:
-        sat_config = config.satellite
 
         sliced_datasets_dict["sat"] = select_time_slice(
             datasets_dict["sat"],
             t0,
-            time_resolution=minutes(sat_config.time_resolution_minutes),
-            interval_start=minutes(sat_config.interval_start_minutes),
-            interval_end=minutes(sat_config.interval_end_minutes),
+            time_resolution=minutes(config.satellite.time_resolution_minutes),
+            interval_start=minutes(config.satellite.interval_start_minutes),
+            interval_end=minutes(config.satellite.interval_end_minutes),
         )
 
     if "generation" in datasets_dict:
-        generation_config = config.generation
         for key, window_config in (
-            ("generation_input", generation_config.input),
-            ("generation_target", generation_config.target),
+            ("generation_input", config.generation.input),
+            ("generation_target", config.generation.target),
         ):
             if window_config is None:
                 continue
@@ -175,7 +172,7 @@ def slice_datasets_by_time(
             sliced_datasets_dict[key] = select_time_slice(
                 datasets_dict["generation"],
                 t0,
-                time_resolution=minutes(generation_config.time_resolution_minutes),
+                time_resolution=minutes(config.generation.time_resolution_minutes),
                 interval_start=minutes(window_config.interval_start_minutes),
                 interval_end=minutes(window_config.interval_end_minutes),
             )
diff --git a/tests/conftest.py b/tests/conftest.py
index db8bfa5a..d2740332 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -192,23 +192,52 @@ def nwp_cloudcasting_zarr_path(session_tmp_path, session_rng):
     yield save_zarr(ds, session_tmp_path, "cloudcasting.zarr", chunks)
 
 
-@pytest.fixture(scope="session")
-def ds_generation(session_rng):
-    times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="30min")
-    location_ids = np.arange(318)
-    # Rough UK bounding box
+def _locations_dataset(session_rng, location_ids):
+    """Build a locations catalog over the given IDs, with random points in a rough UK bbox."""
     lat_min, lat_max = 49.9, 58.7
     lon_min, lon_max = -8.6, 1.8
 
-    # Generate random uniform points
-    longitudes = session_rng.uniform(lon_min, lon_max, len(location_ids)).astype("float64")
-    latitudes = session_rng.uniform(lat_min, lat_max, len(location_ids)).astype("float64")
+    return xr.Dataset(
+        data_vars={
+            "longitude": (
+                "location_id",
+                session_rng.uniform(lon_min, lon_max, len(location_ids)).astype("float64"),
+            ),
+            "latitude": (
+                "location_id",
+                session_rng.uniform(lat_min, lat_max, len(location_ids)).astype("float64"),
+            ),
+        },
+        coords={"location_id": location_ids},
+    )
+
+
+@pytest.fixture(scope="session")
+def ds_locations(session_rng):
+    """The locations catalog - the source of truth for which locations are samplable."""
+    return _locations_dataset(session_rng, np.arange(1, 318))
+
+
+@pytest.fixture(scope="session")
+def ds_site_locations(session_rng):
+    """The locations catalog for the site-level fixtures."""
+    return _locations_dataset(session_rng, np.arange(1, 11))
+
+
+@pytest.fixture(scope="session")
+def ds_generation(session_rng, ds_locations):
+    """Generation for every catalogued location, plus ID 0.
+
+    ID 0 is a summation-model placeholder which the catalog deliberately does not list, so this
+    exercises the filtering of generation IDs down to the catalog's locations.
+    """
+    times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="30min")
+    location_ids = np.concatenate([[0], ds_locations["location_id"].values])
 
     capacity = np.ones((len(times), len(location_ids)))
 
     generation = session_rng.uniform(0, 200, (len(times), len(location_ids))).astype(np.float32)
 
-    # Build Dataset
     return xr.Dataset(
         data_vars={
             "capacity_mwp": (("time_utc", "location_id"), capacity),
@@ -217,29 +246,20 @@ def ds_generation(session_rng):
         coords={
             "time_utc": times,
             "location_id": location_ids,
-            "longitude": ("location_id", longitudes),
-            "latitude": ("location_id", latitudes),
         },
     )
 
 
 # location data (non overlapping time periods) and starting with id 1
 @pytest.fixture(scope="session")
-def ds_site_generation(session_rng):
+def ds_site_generation(session_rng, ds_site_locations):
     # Define a global time range (covers all possible site periods)
     global_times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="30min")
     n_times = len(global_times)
 
-    location_ids = np.arange(1, 11)
+    location_ids = ds_site_locations["location_id"].values
     n_sites = len(location_ids)
 
-    # Rough UK bounding box
-    lat_min, lat_max = 49.9, 58.7
-    lon_min, lon_max = -8.6, 1.8
-
-    longitudes = session_rng.uniform(lon_min, lon_max, n_sites).astype("float64")
-    latitudes = session_rng.uniform(lat_min, lat_max, n_sites).astype("float64")
-
     # Initialize with NaNs
     capacity = np.full((n_times, n_sites), np.nan, dtype="float32")
     generation = np.full((n_times, n_sites), np.nan, dtype="float32")
@@ -258,7 +278,6 @@ def ds_site_generation(session_rng):
             "float32",
         )
 
-    # Build Dataset
     return xr.Dataset(
         data_vars={
             "capacity_mwp": (("time_utc", "location_id"), capacity),
@@ -267,8 +286,6 @@ def ds_site_generation(session_rng):
         coords={
             "time_utc": global_times,
             "location_id": location_ids,
-            "longitude": ("location_id", longitudes),
-            "latitude": ("location_id", latitudes),
         },
     )
 
@@ -283,35 +300,14 @@ def site_generation_zarr_path(session_tmp_path, ds_site_generation):
     yield save_zarr(ds_site_generation, session_tmp_path, "site_generation.zarr")
 
 
-def _locations_dataset_from_generation(ds_generation):
-    """Build a standalone locations metadata dataset from a generation dataset's location coords.
-
-    Excludes location_id 0 - in the generation fixtures that's a placeholder used only for
-    summation models, not a real samplable location, so a properly curated locations catalog
-    wouldn't list it even though generation does.
-    """
-    ds_generation = ds_generation.sel(
-        location_id=[loc_id for loc_id in ds_generation["location_id"].values if loc_id != 0],
-    )
-    return xr.Dataset(
-        data_vars={
-            "longitude": ("location_id", ds_generation["longitude"].values),
-            "latitude": ("location_id", ds_generation["latitude"].values),
-        },
-        coords={"location_id": ds_generation["location_id"].values},
-    )
-
-
 @pytest.fixture(scope="session")
-def locations_zarr_path(session_tmp_path, ds_generation):
-    ds_locations = _locations_dataset_from_generation(ds_generation)
+def locations_zarr_path(session_tmp_path, ds_locations):
     yield save_zarr(ds_locations, session_tmp_path, "locations.zarr")
 
 
 @pytest.fixture(scope="session")
-def site_locations_zarr_path(session_tmp_path, ds_site_generation):
-    ds_locations = _locations_dataset_from_generation(ds_site_generation)
-    yield save_zarr(ds_locations, session_tmp_path, "site_locations.zarr")
+def site_locations_zarr_path(session_tmp_path, ds_site_locations):
+    yield save_zarr(ds_site_locations, session_tmp_path, "site_locations.zarr")
 
 
 @pytest.fixture()
diff --git a/tests/datasets/pvnet/test_preprocess.py b/tests/datasets/pvnet/test_preprocess.py
index f6040c46..74a002c9 100644
--- a/tests/datasets/pvnet/test_preprocess.py
+++ b/tests/datasets/pvnet/test_preprocess.py
@@ -39,20 +39,25 @@ def test_fill_nans_in_dataset_dicts(config_filename):
     assert np.array_equal(datasets_dict["nwp"]["ukv"].values, np.array([-2.0, 3.0, -2.0]))
 
 
-def test_normalise_dataset_dicts_generation():
-    """Generation is normalised to a capacity factor (generation_mw / capacity_mwp)."""
-    generation_mw = np.array([[50.0, 0.0], [100.0, 20.0]])
-    capacity_mwp = np.array([[100.0, 0.0], [100.0, 40.0]])
-    generation = xr.DataArray(
+def _generation_da(generation_mw: np.ndarray, capacity_mwp: np.ndarray) -> xr.DataArray:
+    """Build a generation DataArray from 2D `(time_utc, location_id)` MW and capacity arrays."""
+    n_times, n_locations = generation_mw.shape
+    return xr.DataArray(
         np.stack([generation_mw, capacity_mwp], axis=-1),
         coords={
-            "time_utc": ["2023-01-01T00:00", "2023-01-01T00:30"],
-            "location_id": [1, 2],
+            "time_utc": pd.date_range("2023-01-01", periods=n_times, freq="30min"),
+            "location_id": np.arange(1, n_locations + 1),
             "gen_param": ["generation_mw", "capacity_mwp"],
         },
         dims=("time_utc", "location_id", "gen_param"),
     )
 
+
+def test_normalise_dataset_dicts_generation():
+    """Generation is normalised to a capacity factor (generation_mw / capacity_mwp)."""
+    capacity_mwp = np.array([[100.0, 0.0], [100.0, 40.0]])
+    generation = _generation_da(np.array([[50.0, 0.0], [100.0, 20.0]]), capacity_mwp)
+
     datasets_dict = {
         "generation_input": generation,
         "generation_target": generation.copy(deep=True),
@@ -76,56 +81,6 @@ def test_normalise_dataset_dicts_generation():
         assert np.array_equal(result.sel(gen_param="capacity_mwp").values, capacity_mwp)
 
 
-def _generation_source(generation_mw: np.ndarray, capacity_mwp: np.ndarray) -> xr.DataArray:
-    """Build a generation DataArray with `n` timesteps for a single location."""
-    return xr.DataArray(
-        np.stack([generation_mw[:, None], capacity_mwp[:, None]], axis=-1),
-        coords={
-            "time_utc": pd.date_range("2023-01-01", periods=len(generation_mw), freq="30min"),
-            "location_id": [1],
-            "gen_param": ["generation_mw", "capacity_mwp"],
-        },
-        dims=("time_utc", "location_id", "gen_param"),
-    )
-
-
-def test_normalise_does_not_mutate_shared_source():
-    """Normalising windows must not write through to the array they were sliced from.
-
-    Time slices are views onto the eagerly-loaded source, so an in-place write would corrupt it
-    for every later sample and double-normalise timesteps shared by the two windows.
-    """
-    src = _generation_source(np.array([100.0, 150.0, 200.0]), np.full(3, 100.0))
-    before = src.values.copy()
-
-    # Windows deliberately overlap on the middle timestep
-    datasets_dict = {
-        "generation_input": src.isel(time_utc=slice(0, 2)),
-        "generation_target": src.isel(time_utc=slice(1, 3)),
-    }
-    datasets_dict = normalise_dataset_dicts(datasets_dict, {}, {}, {}, {})
-
-    assert np.array_equal(src.values, before), "source array was mutated by normalisation"
-
-    def generation_of(key: str) -> np.ndarray:
-        return datasets_dict[key].sel(gen_param="generation_mw").values.ravel()
-
-    assert np.array_equal(generation_of("generation_input"), [1.0, 1.5])
-    # The shared timestep is normalised once, not once per window
-    assert np.array_equal(generation_of("generation_target"), [1.5, 2.0])
-
-
-def test_normalise_generation_is_dimension_order_agnostic():
-    """`gen_param` is indexed by name, so it need not be the last dimension."""
-    src = _generation_source(np.array([100.0, 150.0, 200.0]), np.full(3, 100.0))
-    transposed = src.transpose("gen_param", "time_utc", "location_id")
-
-    result = normalise_dataset_dicts({"generation_input": transposed}, {}, {}, {}, {})
-
-    normalised = result["generation_input"].sel(gen_param="generation_mw").values.ravel()
-    assert np.array_equal(normalised, [1.0, 1.5, 2.0])
-
-
 def test_apply_dropout_to_datasets(pvnet_config_filename):
     config = load_yaml_configuration(pvnet_config_filename)
 
diff --git a/tests/load/test_load_locations.py b/tests/load/test_load_locations.py
index e4bed3cb..38de3a09 100644
--- a/tests/load/test_load_locations.py
+++ b/tests/load/test_load_locations.py
@@ -34,23 +34,3 @@ def test_open_locations_missing_data_var(tmp_path: Path):
 
     with pytest.raises(ValueError, match="Locations data should have variables"):
         open_locations(zarr_path=str(zarr_path))
-
-
-def test_open_locations_bad_dtype(tmp_path: Path):
-    """Test that open_locations raises a TypeError on incorrect data dtypes."""
-    zarr_path = tmp_path / "bad_locations.zarr"
-
-    # Create dataset where longitude is integer
-    bad_ds = xr.Dataset(
-        data_vars={
-            "longitude": (("location_id",), [0, 1]),
-            "latitude": (("location_id",), [0.0, 1.0]),
-        },
-        coords={
-            "location_id": [1, 2],
-        },
-    )
-    bad_ds.to_zarr(zarr_path)
-
-    with pytest.raises(TypeError, match="longitude in locations data should be floating"):
-        open_locations(zarr_path=str(zarr_path))

From e2323bb0a4d96cc95688218a48768191eb78ef76 Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Fri, 31 Jul 2026 08:55:41 +0000
Subject: [PATCH 4/9] Move global variable to its rightful place at the top

---
 tests/config/test_config.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tests/config/test_config.py b/tests/config/test_config.py
index d0e31714..efca2716 100644
--- a/tests/config/test_config.py
+++ b/tests/config/test_config.py
@@ -4,6 +4,8 @@
 from ocf_data_sampler.config.load import load_yaml_configuration
 from ocf_data_sampler.config.model import PVNetDataConfig
 
+_MINIMAL_SAMPLING_GRID = {"locations_zarr_path": "locations.zarr", "t0_resolution_minutes": 30}
+
 
 def _load_config_and_provider(config_path):
     config = load_yaml_configuration(config_path)
@@ -16,9 +18,6 @@ def _validate_configuration(config):
     return PVNetDataConfig(**config.model_dump())
 
 
-_MINIMAL_SAMPLING_GRID = {"locations_zarr_path": "locations.zarr", "t0_resolution_minutes": 30}
-
-
 def test_default_configuration():
     """Test default pydantic class - sampling_grid is the only required field"""
     _ = PVNetDataConfig(sampling_grid=_MINIMAL_SAMPLING_GRID)

From e25fc68b931c1e97c14e5e403a88dc7dfdc08ebb Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Fri, 31 Jul 2026 15:28:02 +0000
Subject: [PATCH 5/9] Convert locations input to csv and add option to exclude
 IDs from config

---
 src/ocf_data_sampler/config/model.py          |  21 +++-
 .../datasets/pvnet/dataset.py                 |  48 ++++----
 src/ocf_data_sampler/load/locations.py        |  60 +++++-----
 tests/config/test_config.py                   |   2 +-
 tests/conftest.py                             |  57 +++++-----
 tests/datasets/pvnet/test_dataset.py          |  19 ++++
 tests/fixtures/configs/pvnet_test_config.yaml |   2 +-
 tests/load/test_load_locations.py             | 106 +++++++++++++-----
 8 files changed, 209 insertions(+), 106 deletions(-)

diff --git a/src/ocf_data_sampler/config/model.py b/src/ocf_data_sampler/config/model.py
index 5ae27cc7..ab2e37f4 100644
--- a/src/ocf_data_sampler/config/model.py
+++ b/src/ocf_data_sampler/config/model.py
@@ -394,11 +394,16 @@ def validate_windows(self) -> "Generation":
 class SamplingGrid(Base):
     """Configuration for the (location, time) grid that t0 times are sampled from."""
 
-    locations_zarr_path: str = Field(
+    locations_csv_path: str = Field(
         ...,
-        description="Absolute or relative zarr filepath to the locations metadata (location IDs "
-        "and their coordinates) - see `ocf_data_sampler.load.locations.open_locations`. Prefix "
-        "with a protocol like s3:// to read from alternative filesystems.",
+        description="Absolute or relative CSV filepath to the locations metadata (location IDs "
+        "and their coordinates) - see `ocf_data_sampler.load.locations.open_locations`.",
+    )
+
+    exclude_location_ids: list[int] = Field(
+        default=[],
+        description="Location IDs from the locations metadata to drop from the sampling grid. "
+        "Every ID listed must be present in the locations data.",
     )
 
     t0_resolution_minutes: int = Field(
@@ -408,6 +413,14 @@ class SamplingGrid(Base):
         "times regardless of which other input sources are configured.",
     )
 
+    @field_validator("exclude_location_ids")
+    def validate_exclude_location_ids_unique(cls, v: list[int]) -> list[int]:
+        """Validate 'exclude_location_ids'."""
+        duplicates = {i for i in v if v.count(i) > 1}
+        if duplicates:
+            raise ValueError(f"exclude_location_ids contains duplicates: {sorted(duplicates)}")
+        return v
+
 
 class SolarPosition(TimeWindowMixin):
     """Solar position configuration model."""
diff --git a/src/ocf_data_sampler/datasets/pvnet/dataset.py b/src/ocf_data_sampler/datasets/pvnet/dataset.py
index 8fdedf9a..94f59a8e 100644
--- a/src/ocf_data_sampler/datasets/pvnet/dataset.py
+++ b/src/ocf_data_sampler/datasets/pvnet/dataset.py
@@ -45,29 +45,36 @@
 
 
 
-def get_locations(zarr_path: str) -> list[Location]:
+def get_locations(csv_path: str, exclude_ids: list[int] | None = None) -> list[Location]:
     """Load the locations metadata and build the list of all locations.
 
     Args:
-        zarr_path: Path to the locations zarr data
+        csv_path: Path to the locations CSV data
+        exclude_ids: Location IDs to drop from the returned locations
     """
-    locations_data = open_locations(zarr_path)
-
-    locations = []
-    location_ids = locations_data["location_id"].values
-
-    for location_id in location_ids:
-        loc_data = locations_data.sel(location_id=location_id)
-        locations.append(
-            Location(
-                x=loc_data["longitude"].values,
-                y=loc_data["latitude"].values,
-                coord_system="lon_lat",
-                id=int(location_id),
-            ),
-        )
+    locations_data = open_locations(csv_path)
 
-    return locations
+    if exclude_ids:
+        missing_ids = np.setdiff1d(exclude_ids, locations_data["location_id"].values)
+        if len(missing_ids) > 0:
+            raise ValueError(
+                f"Cannot exclude location IDs which are not in the locations data: {missing_ids}",
+            )
+
+        locations_data = locations_data[~locations_data["location_id"].isin(exclude_ids)]
+
+        if len(locations_data) == 0:
+            raise ValueError("All location IDs in the locations data have been excluded")
+
+    return [
+        Location(
+            x=row.longitude,
+            y=row.latitude,
+            coord_system="lon_lat",
+            id=int(row.location_id),
+        )
+        for row in locations_data.itertuples()
+    ]
 
 
 def xarray_to_lightarray_dict(
@@ -255,7 +262,10 @@ def __init__(
 
         config = load_yaml_configuration(config_filename)
 
-        locations = get_locations(config.sampling_grid.locations_zarr_path)
+        locations = get_locations(
+            config.sampling_grid.locations_csv_path,
+            config.sampling_grid.exclude_location_ids,
+        )
 
         datasets_dict = get_dataset_dict(config)
 
diff --git a/src/ocf_data_sampler/load/locations.py b/src/ocf_data_sampler/load/locations.py
index d888fac5..03112e92 100644
--- a/src/ocf_data_sampler/load/locations.py
+++ b/src/ocf_data_sampler/load/locations.py
@@ -1,47 +1,51 @@
 """Functions for loading locations metadata.
 
-Locations data schema: a Zarr file with the following data variables and dimensions/coordinates:
-
-Dimensions: (location_id,)
-Data Variables:
-    longitude (location_id): The longitudes of the locations
-    latitude (location_id): The latitudes of the locations
-Coordinates:
-    location_id (location_id): The integer IDs of the locations
+Locations data schema: a CSV file with the following columns:
+
+    location_id: The integer IDs of the locations
+    longitude: The longitudes of the locations
+    latitude: The latitudes of the locations
+
+Rows must be in increasing `location_id` order, and no value may be left blank (i.e. no NaNs).
+
+A CSV is used rather than zarr since this catalogue is small and benefits from being human
+readable and hand editable. Additional columns may be included to make the file easier to work with,
+but are dropped on load.
 """
 
-import numpy as np
-import xarray as xr
+import pandas as pd
 
 from ocf_data_sampler.common.indexing import assert_values_unique_increasing
-from ocf_data_sampler.load.conventions import validate_coords
 
 
-def open_locations(zarr_path: str) -> xr.Dataset:
-    """Open and eagerly load the locations metadata and validate its data types.
+def open_locations(csv_path: str) -> pd.DataFrame:
+    """Open the locations metadata and validate its columns and data types.
 
     Args:
-        zarr_path: Path to the locations zarr data
+        csv_path: Path to the locations CSV data
 
     Returns:
-        xr.Dataset: The opened locations metadata
+        pd.DataFrame: The locations metadata, in increasing location ID order
     """
-    ds = xr.open_zarr(zarr_path, chunks=None)
-
-    if set(ds.data_vars) != {"longitude", "latitude"}:
+    column_dtypes = {
+        "location_id": "int64",
+        "longitude": "float64",
+        "latitude": "float64",
+    }
+    df = pd.read_csv(csv_path, dtype=column_dtypes)
+
+    if missing_columns := set(column_dtypes) - set(df.columns):
         raise ValueError(
-            f"Locations data should have variables 'longitude' and 'latitude', "
-            f"but found {set(ds.data_vars)} instead."
+            f"Locations data should have columns {list(column_dtypes)}, but the following "
+            f"were missing: {missing_columns}",
         )
 
-    validate_coords(ds, {"location_id": np.integer}, source="locations data")
-
-    for var in ("longitude", "latitude"):
-        if not np.issubdtype(ds[var].dtype, np.floating):
-            raise TypeError(f"{var} in locations data should be floating, not {ds[var].dtype}")
+    # Extra columns are allowed in the file, but are dropped on load to avoid bloat
+    df = df[list(column_dtypes)]
 
-    ds = ds.load()
+    if df.isna().any().any():
+        raise ValueError("Locations data must not contain missing values")
 
-    assert_values_unique_increasing(ds["location_id"].values, "location_id")
+    assert_values_unique_increasing(df["location_id"].values, "location_id")
 
-    return ds
+    return df
diff --git a/tests/config/test_config.py b/tests/config/test_config.py
index efca2716..4b7f1719 100644
--- a/tests/config/test_config.py
+++ b/tests/config/test_config.py
@@ -4,7 +4,7 @@
 from ocf_data_sampler.config.load import load_yaml_configuration
 from ocf_data_sampler.config.model import PVNetDataConfig
 
-_MINIMAL_SAMPLING_GRID = {"locations_zarr_path": "locations.zarr", "t0_resolution_minutes": 30}
+_MINIMAL_SAMPLING_GRID = {"locations_csv_path": "locations.csv", "t0_resolution_minutes": 30}
 
 
 def _load_config_and_provider(config_path):
diff --git a/tests/conftest.py b/tests/conftest.py
index d2740332..9708b0c9 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -62,6 +62,13 @@ def create_xr_dataset(coords, data, name, attrs=None):
     return da.to_dataset(name=name)
 
 
+def save_csv(df, path, filename):
+    """Save dataframe to csv"""
+    csv_path = path / filename
+    df.to_csv(csv_path, index=False)
+    return str(csv_path)
+
+
 def save_zarr(ds, path, filename, chunks=None):
     """Save dataset to zarr"""
     if chunks:
@@ -192,47 +199,41 @@ def nwp_cloudcasting_zarr_path(session_tmp_path, session_rng):
     yield save_zarr(ds, session_tmp_path, "cloudcasting.zarr", chunks)
 
 
-def _locations_dataset(session_rng, location_ids):
+def _locations_dataframe(session_rng, location_ids):
     """Build a locations catalog over the given IDs, with random points in a rough UK bbox."""
     lat_min, lat_max = 49.9, 58.7
     lon_min, lon_max = -8.6, 1.8
 
-    return xr.Dataset(
-        data_vars={
-            "longitude": (
-                "location_id",
-                session_rng.uniform(lon_min, lon_max, len(location_ids)).astype("float64"),
-            ),
-            "latitude": (
-                "location_id",
-                session_rng.uniform(lat_min, lat_max, len(location_ids)).astype("float64"),
-            ),
+    return pd.DataFrame(
+        {
+            "location_id": location_ids,
+            "longitude": session_rng.uniform(lon_min, lon_max, len(location_ids)),
+            "latitude": session_rng.uniform(lat_min, lat_max, len(location_ids)),
         },
-        coords={"location_id": location_ids},
     )
 
 
 @pytest.fixture(scope="session")
-def ds_locations(session_rng):
+def df_locations(session_rng):
     """The locations catalog - the source of truth for which locations are samplable."""
-    return _locations_dataset(session_rng, np.arange(1, 318))
+    return _locations_dataframe(session_rng, np.arange(1, 318))
 
 
 @pytest.fixture(scope="session")
-def ds_site_locations(session_rng):
+def df_site_locations(session_rng):
     """The locations catalog for the site-level fixtures."""
-    return _locations_dataset(session_rng, np.arange(1, 11))
+    return _locations_dataframe(session_rng, np.arange(1, 11))
 
 
 @pytest.fixture(scope="session")
-def ds_generation(session_rng, ds_locations):
+def ds_generation(session_rng, df_locations):
     """Generation for every catalogued location, plus ID 0.
 
     ID 0 is a summation-model placeholder which the catalog deliberately does not list, so this
     exercises the filtering of generation IDs down to the catalog's locations.
     """
     times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="30min")
-    location_ids = np.concatenate([[0], ds_locations["location_id"].values])
+    location_ids = np.concatenate([[0], df_locations["location_id"].values])
 
     capacity = np.ones((len(times), len(location_ids)))
 
@@ -252,12 +253,12 @@ def ds_generation(session_rng, ds_locations):
 
 # location data (non overlapping time periods) and starting with id 1
 @pytest.fixture(scope="session")
-def ds_site_generation(session_rng, ds_site_locations):
+def ds_site_generation(session_rng, df_site_locations):
     # Define a global time range (covers all possible site periods)
     global_times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="30min")
     n_times = len(global_times)
 
-    location_ids = ds_site_locations["location_id"].values
+    location_ids = df_site_locations["location_id"].values
     n_sites = len(location_ids)
 
     # Initialize with NaNs
@@ -301,13 +302,13 @@ def site_generation_zarr_path(session_tmp_path, ds_site_generation):
 
 
 @pytest.fixture(scope="session")
-def locations_zarr_path(session_tmp_path, ds_locations):
-    yield save_zarr(ds_locations, session_tmp_path, "locations.zarr")
+def locations_csv_path(session_tmp_path, df_locations):
+    yield save_csv(df_locations, session_tmp_path, "locations.csv")
 
 
 @pytest.fixture(scope="session")
-def site_locations_zarr_path(session_tmp_path, ds_site_locations):
-    yield save_zarr(ds_site_locations, session_tmp_path, "site_locations.zarr")
+def site_locations_csv_path(session_tmp_path, df_site_locations):
+    yield save_csv(df_site_locations, session_tmp_path, "site_locations.csv")
 
 
 @pytest.fixture()
@@ -316,14 +317,14 @@ def pvnet_config_filename(
     config_filename,
     nwp_ukv_zarr_path,
     generation_zarr_path,
-    locations_zarr_path,
+    locations_csv_path,
     sat_zarr_path,
 ):
     config = load_yaml_configuration(config_filename)
     config.nwp["ukv"].zarr_path = nwp_ukv_zarr_path
     config.satellite.zarr_path = sat_zarr_path
     config.generation.zarr_path = generation_zarr_path
-    config.sampling_grid.locations_zarr_path = locations_zarr_path
+    config.sampling_grid.locations_csv_path = locations_csv_path
 
     path = tmp_path / "configuration.yaml"
     save_yaml_configuration(config, str(path))
@@ -336,14 +337,14 @@ def pvnet_site_config_filename(
     config_filename,
     nwp_ukv_zarr_path,
     site_generation_zarr_path,
-    site_locations_zarr_path,
+    site_locations_csv_path,
     sat_zarr_path,
 ):
     config = load_yaml_configuration(config_filename)
     config.nwp["ukv"].zarr_path = nwp_ukv_zarr_path
     config.satellite.zarr_path = sat_zarr_path
     config.generation.zarr_path = site_generation_zarr_path
-    config.sampling_grid.locations_zarr_path = site_locations_zarr_path
+    config.sampling_grid.locations_csv_path = site_locations_csv_path
 
     path = session_tmp_path / "configuration.yaml"
     save_yaml_configuration(config, str(path))
diff --git a/tests/datasets/pvnet/test_dataset.py b/tests/datasets/pvnet/test_dataset.py
index b3d64f6c..8ad96b16 100644
--- a/tests/datasets/pvnet/test_dataset.py
+++ b/tests/datasets/pvnet/test_dataset.py
@@ -11,6 +11,7 @@
 from ocf_data_sampler.datasets.pvnet.dataset import (
     PVNetConcurrentDataset,
     PVNetDataset,
+    get_locations,
     get_time_periods_mask,
 )
 
@@ -135,6 +136,24 @@ def test_pvnet_dataset(pvnet_config_filename):
     _pvnet_dataset_sample_check(sample, dataset.config)
 
 
+def test_get_locations_exclude_ids(locations_csv_path):
+    excluded_ids = [1, 5, 317]
+    locations = get_locations(locations_csv_path, exclude_ids=excluded_ids)
+
+    assert len(locations) == 317 - len(excluded_ids)
+    assert not set(excluded_ids) & {loc.id for loc in locations}
+
+
+def test_get_locations_exclude_unknown_id(locations_csv_path):
+    with pytest.raises(ValueError, match="not in the locations data"):
+        get_locations(locations_csv_path, exclude_ids=[1, 9999])
+
+
+def test_get_locations_exclude_all_ids(locations_csv_path):
+    with pytest.raises(ValueError, match=r"All location IDs .* have been excluded"):
+        get_locations(locations_csv_path, exclude_ids=list(range(1, 318)))
+
+
 def test_pvnet_dataset_sites(pvnet_site_config_filename):
     dataset = PVNetDataset(
         pvnet_site_config_filename,
diff --git a/tests/fixtures/configs/pvnet_test_config.yaml b/tests/fixtures/configs/pvnet_test_config.yaml
index cf868500..b1f1dfd3 100644
--- a/tests/fixtures/configs/pvnet_test_config.yaml
+++ b/tests/fixtures/configs/pvnet_test_config.yaml
@@ -1,5 +1,5 @@
 sampling_grid:
-  locations_zarr_path: set_in_temp_file
+  locations_csv_path: set_in_temp_file
   t0_resolution_minutes: 30
 
 generation:
diff --git a/tests/load/test_load_locations.py b/tests/load/test_load_locations.py
index 38de3a09..197e0bb5 100644
--- a/tests/load/test_load_locations.py
+++ b/tests/load/test_load_locations.py
@@ -1,36 +1,92 @@
 from pathlib import Path
 
 import numpy as np
+import pandas as pd
 import pytest
-import xarray as xr
 
 from ocf_data_sampler.load.locations import open_locations
 
 
-def test_open_locations(locations_zarr_path):
+def _write_csv(tmp_path: Path, df: pd.DataFrame) -> str:
+    csv_path = tmp_path / "locations.csv"
+    df.to_csv(csv_path, index=False)
+    return str(csv_path)
+
+
+def test_open_locations(locations_csv_path):
     """Test the locations data loader with valid data."""
-    ds = open_locations(locations_zarr_path)
-
-    assert isinstance(ds, xr.Dataset)
-    assert set(ds.data_vars) == {"longitude", "latitude"}
-    assert ds["longitude"].dims == ("location_id",)
-    assert ds["latitude"].dims == ("location_id",)
-    assert len(np.unique(ds.coords["location_id"])) == ds.sizes["location_id"]
-
-
-def test_open_locations_missing_data_var(tmp_path: Path):
-    """Test that open_locations raises a ValueError when a required data variable is missing."""
-    zarr_path = tmp_path / "bad_locations.zarr"
-
-    bad_ds = xr.Dataset(
-        data_vars={
-            "longitude": (("location_id",), [0.0, 1.0]),
-        },
-        coords={
-            "location_id": [1, 2],
-        },
+    df = open_locations(locations_csv_path)
+
+    assert isinstance(df, pd.DataFrame)
+    assert list(df.columns) == ["location_id", "longitude", "latitude"]
+    assert df["location_id"].dtype == "int64"
+    assert df["longitude"].dtype == "float64"
+    assert df["latitude"].dtype == "float64"
+    assert df["location_id"].is_unique
+
+
+def test_open_locations_unsorted_ids(tmp_path: Path):
+    """The file must be kept in increasing location ID order."""
+    csv_path = _write_csv(
+        tmp_path,
+        pd.DataFrame(
+            {
+                "location_id": [7, 2, 5],
+                "longitude": [1.0, 0.0, 0.5],
+                "latitude": [51.0, 50.0, 50.5],
+            },
+        ),
+    )
+
+    with pytest.raises(ValueError, match="location_id must be strictly increasing"):
+        open_locations(csv_path)
+
+
+def test_open_locations_missing_value(tmp_path: Path):
+    """A row with a blank coordinate is rejected rather than loaded as NaN."""
+    csv_path = _write_csv(
+        tmp_path,
+        pd.DataFrame({"location_id": [1], "longitude": [0.0], "latitude": [np.nan]}),
+    )
+
+    with pytest.raises(ValueError, match="must not contain missing values"):
+        open_locations(csv_path)
+
+
+def test_open_locations_missing_column(tmp_path: Path):
+    """Test that open_locations raises a ValueError when a required column is missing."""
+    csv_path = _write_csv(
+        tmp_path,
+        pd.DataFrame({"location_id": [1, 2], "longitude": [0.0, 1.0]}),
+    )
+
+    with pytest.raises(ValueError, match="Locations data should have columns"):
+        open_locations(csv_path)
+
+
+def test_open_locations_duplicate_ids(tmp_path: Path):
+    """Test that open_locations raises a ValueError on repeated location IDs."""
+    csv_path = _write_csv(
+        tmp_path,
+        pd.DataFrame(
+            {
+                "location_id": [1, 1],
+                "longitude": [0.0, 1.0],
+                "latitude": [50.0, 51.0],
+            },
+        ),
+    )
+
+    with pytest.raises(ValueError, match="location_id must be strictly increasing"):
+        open_locations(csv_path)
+
+
+def test_open_locations_non_integer_id(tmp_path: Path):
+    """Test that open_locations rejects location IDs which are not integers."""
+    csv_path = _write_csv(
+        tmp_path,
+        pd.DataFrame({"location_id": [1.5], "longitude": [0.0], "latitude": [50.0]}),
     )
-    bad_ds.to_zarr(zarr_path)
 
-    with pytest.raises(ValueError, match="Locations data should have variables"):
-        open_locations(zarr_path=str(zarr_path))
+    with pytest.raises(ValueError, match="int64"):
+        open_locations(csv_path)

From 045fc29ba1624ccd44ce084a7c5b41796196e48f Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Mon, 3 Aug 2026 08:10:31 +0000
Subject: [PATCH 6/9] Add input/target to eror message

---
 src/ocf_data_sampler/datasets/pvnet/valid_t0s.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py b/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py
index a824fc87..7921bc97 100644
--- a/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py
+++ b/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py
@@ -81,7 +81,10 @@ def find_valid_time_periods(
             raise ValueError("No valid t0 periods found for satellite data")
 
     if "generation" in datasets_dict:
-        for window_config in (config.generation.input, config.generation.target):
+        for window_name, window_config in (
+            ("input", config.generation.input),
+            ("target", config.generation.target),
+        ):
             if window_config is None:
                 continue
 
@@ -93,14 +96,15 @@ def find_valid_time_periods(
             )
 
             if len(time_periods) == 0:
-                raise ValueError("No valid t0 periods found for generation data")
+                raise ValueError(
+                    f"No valid t0 periods found for {window_name} generation data",
+                )
 
             contiguous_time_periods.append(time_periods)
 
     # Find joint overlapping contiguous time periods
     valid_time_periods = intersect_time_periods(contiguous_time_periods)
 
-
     # check there are some valid time periods
     if len(valid_time_periods) == 0:
         raise ValueError(f"No valid time periods found, {contiguous_time_periods=}")

From 2c35348a857b16b3cae51c4b7ec55c0eb33d9504 Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Wed, 5 Aug 2026 10:56:27 +0000
Subject: [PATCH 7/9] Fix location IDs in tests

---
 tests/conftest.py                             | 30 +++++++++++--------
 tests/datasets/pvnet/test_dataset.py          | 27 ++++++++++-------
 tests/fixtures/configs/pvnet_test_config.yaml |  1 +
 tests/load/test_load_generation.py            |  4 ++-
 4 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 9708b0c9..685e242a 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -14,6 +14,12 @@
 NWP_FREQ = pd.Timedelta("3h")
 RANDOM_SEED = 42
 
+# The LOCATION_IDS catalog mirrors the GSPs: ID 0 is the national aggregate and IDs 1-317 are the
+# regional GSPs
+LOCATION_IDS = tuple(range(318))
+# The SITE_LOCATION_IDS catalog has no national aggregate since they mirror different sites
+SITE_LOCATION_IDS = tuple(range(1, 11))
+
 UK_SAT_AREA = """msg_seviri_rss_3km:
     description: MSG SEVIRI Rapid Scanning Service area definition with 3 km resolution
     projection:
@@ -216,33 +222,31 @@ def _locations_dataframe(session_rng, location_ids):
 @pytest.fixture(scope="session")
 def df_locations(session_rng):
     """The locations catalog - the source of truth for which locations are samplable."""
-    return _locations_dataframe(session_rng, np.arange(1, 318))
+    return _locations_dataframe(session_rng, LOCATION_IDS)
 
 
 @pytest.fixture(scope="session")
 def df_site_locations(session_rng):
     """The locations catalog for the site-level fixtures."""
-    return _locations_dataframe(session_rng, np.arange(1, 11))
+    return _locations_dataframe(session_rng, SITE_LOCATION_IDS)
 
 
 @pytest.fixture(scope="session")
 def ds_generation(session_rng, df_locations):
-    """Generation for every catalogued location, plus ID 0.
+    """Generation for every catalogued location with no missing generation data"""
 
-    ID 0 is a summation-model placeholder which the catalog deliberately does not list, so this
-    exercises the filtering of generation IDs down to the catalog's locations.
-    """
     times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="30min")
-    location_ids = np.concatenate([[0], df_locations["location_id"].values])
-
-    capacity = np.ones((len(times), len(location_ids)))
+    location_ids = df_locations["location_id"].values
+    shape = (len(times), len(location_ids))
 
-    generation = session_rng.uniform(0, 200, (len(times), len(location_ids))).astype(np.float32)
+    capacity = 200
+    capacities = np.full(shape, fill_value=capacity, dtype="float32")
+    generations = session_rng.uniform(0, capacity, shape).astype("float32")
 
     return xr.Dataset(
         data_vars={
-            "capacity_mwp": (("time_utc", "location_id"), capacity),
-            "generation_mw": (("time_utc", "location_id"), generation),
+            "capacity_mwp": (("time_utc", "location_id"), capacities),
+            "generation_mw": (("time_utc", "location_id"), generations),
         },
         coords={
             "time_utc": times,
@@ -345,6 +349,8 @@ def pvnet_site_config_filename(
     config.satellite.zarr_path = sat_zarr_path
     config.generation.zarr_path = site_generation_zarr_path
     config.sampling_grid.locations_csv_path = site_locations_csv_path
+    # The site catalog has no national aggregate, so nothing to exclude
+    config.sampling_grid.exclude_location_ids = []
 
     path = session_tmp_path / "configuration.yaml"
     save_yaml_configuration(config, str(path))
diff --git a/tests/datasets/pvnet/test_dataset.py b/tests/datasets/pvnet/test_dataset.py
index 8ad96b16..53694def 100644
--- a/tests/datasets/pvnet/test_dataset.py
+++ b/tests/datasets/pvnet/test_dataset.py
@@ -14,6 +14,7 @@
     get_locations,
     get_time_periods_mask,
 )
+from tests.conftest import LOCATION_IDS, SITE_LOCATION_IDS
 
 
 def _pvnet_dataset_sample_check(sample, config, batch_dim = None):
@@ -115,6 +116,11 @@ def test_get_time_periods_mask():
     assert np.array_equal(mask, expected_mask), f"Expected {expected_mask} but got {mask}"
 
 
+def _expected_num_locations(dataset, catalog_ids):
+    """The catalogued locations which survive the config's exclusion list."""
+    return len(catalog_ids) - len(dataset.config.sampling_grid.exclude_location_ids)
+
+
 def test_pvnet_dataset(pvnet_config_filename):
     dataset = PVNetDataset(
         pvnet_config_filename,
@@ -125,7 +131,7 @@ def test_pvnet_dataset(pvnet_config_filename):
     )
 
     expected_t0s = 6  # 2 time periods each with 3 t0s (inclusive) at 30 minute intervals
-    num_locs = 317 # Quantity of regional GSPs
+    num_locs = _expected_num_locations(dataset, LOCATION_IDS)
     assert len(dataset.locations) == num_locs
 
     assert len(dataset.valid_t0_times) == expected_t0s
@@ -137,21 +143,22 @@ def test_pvnet_dataset(pvnet_config_filename):
 
 
 def test_get_locations_exclude_ids(locations_csv_path):
-    excluded_ids = [1, 5, 317]
+    excluded_ids = [LOCATION_IDS[0], LOCATION_IDS[5], LOCATION_IDS[-1]]
     locations = get_locations(locations_csv_path, exclude_ids=excluded_ids)
 
-    assert len(locations) == 317 - len(excluded_ids)
+    assert len(locations) == len(LOCATION_IDS) - len(excluded_ids)
     assert not set(excluded_ids) & {loc.id for loc in locations}
 
 
 def test_get_locations_exclude_unknown_id(locations_csv_path):
+    unknown_id = max(LOCATION_IDS) + 1
     with pytest.raises(ValueError, match="not in the locations data"):
-        get_locations(locations_csv_path, exclude_ids=[1, 9999])
+        get_locations(locations_csv_path, exclude_ids=[unknown_id])
 
 
 def test_get_locations_exclude_all_ids(locations_csv_path):
     with pytest.raises(ValueError, match=r"All location IDs .* have been excluded"):
-        get_locations(locations_csv_path, exclude_ids=list(range(1, 318)))
+        get_locations(locations_csv_path, exclude_ids=list(LOCATION_IDS))
 
 
 def test_pvnet_dataset_sites(pvnet_site_config_filename):
@@ -164,7 +171,7 @@ def test_pvnet_dataset_sites(pvnet_site_config_filename):
     )
 
     expected_t0s = 6  # 2 time periods each with 3 t0s (inclusive) at 30 minute intervals
-    num_locs = 10
+    num_locs = _expected_num_locations(dataset, SITE_LOCATION_IDS)
     assert len(dataset.locations) == num_locs
     # Should be less than num_locs * expected_t0s as not all locations have data for all t0s
     # in the time periods
@@ -195,14 +202,14 @@ def check_samples_equal(sample0, sample1):
 def test_pvnet_concurrent_dataset(pvnet_config_filename):
     # Create dataset object using limited set of GSPs
     dataset = PVNetConcurrentDataset(pvnet_config_filename)
-    num_gsps = 317
-    assert len(dataset.locations) == num_gsps  # Quantity of regional GSPs
+    num_locations = _expected_num_locations(dataset, LOCATION_IDS)
+    assert len(dataset.locations) == num_locations
     # NB. I have not checked the value (39 below) is in fact correct
     assert len(dataset.valid_t0_times) == 39
     assert len(dataset) == 39
 
     sample = dataset[0]
-    _pvnet_dataset_sample_check(sample, dataset.config, (num_gsps,))
+    _pvnet_dataset_sample_check(sample, dataset.config, (num_locations,))
 
 
 def test_pvnet_dataset_getitem_bounds(pvnet_config_filename):
@@ -271,7 +278,7 @@ def test_pvnet_dataset_without_generation(tmp_path, pvnet_config_filename):
     assert dataset.complete_generation
 
     # All locations from the locations catalog are available - none to filter out
-    assert len(dataset.locations) == 317
+    assert len(dataset.locations) == _expected_num_locations(dataset, LOCATION_IDS)
 
     sample = dataset[0]
     assert "generation_input" not in sample
diff --git a/tests/fixtures/configs/pvnet_test_config.yaml b/tests/fixtures/configs/pvnet_test_config.yaml
index b1f1dfd3..34a398e3 100644
--- a/tests/fixtures/configs/pvnet_test_config.yaml
+++ b/tests/fixtures/configs/pvnet_test_config.yaml
@@ -1,6 +1,7 @@
 sampling_grid:
   locations_csv_path: set_in_temp_file
   t0_resolution_minutes: 30
+  exclude_location_ids: [0]
 
 generation:
   zarr_path: set_in_temp_file
diff --git a/tests/load/test_load_generation.py b/tests/load/test_load_generation.py
index 56c5dc04..83d6439e 100644
--- a/tests/load/test_load_generation.py
+++ b/tests/load/test_load_generation.py
@@ -6,6 +6,7 @@
 import xarray as xr
 
 from ocf_data_sampler.load.generation import open_generation
+from tests.conftest import LOCATION_IDS
 
 
 def test_open_generation(generation_zarr_path):
@@ -14,7 +15,8 @@ def test_open_generation(generation_zarr_path):
 
     assert isinstance(da, xr.DataArray)
     assert da.dims == ("time_utc", "location_id", "gen_param")
-    assert da.shape == (49, 318, 2)
+    # 24 hours of 30 minute data (inclusive), every catalogued location, capacity + generation
+    assert da.shape == (49, len(LOCATION_IDS), 2)
     assert len(np.unique(da.coords["location_id"])) == da.shape[1]
 
 

From 834c1f2f85f8b73ec1f1654389ee52a03a580cb9 Mon Sep 17 00:00:00 2001
From: James Fulton <41546094+dfulu@users.noreply.github.com>
Date: Thu, 6 Aug 2026 10:01:56 +0100
Subject: [PATCH 8/9] Update src/ocf_data_sampler/datasets/pvnet/preprocess.py

Co-authored-by: Alexandra Udaltsova <43303448+AUdaltsova@users.noreply.github.com>
---
 src/ocf_data_sampler/datasets/pvnet/preprocess.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ocf_data_sampler/datasets/pvnet/preprocess.py b/src/ocf_data_sampler/datasets/pvnet/preprocess.py
index ee61b54e..144b374f 100644
--- a/src/ocf_data_sampler/datasets/pvnet/preprocess.py
+++ b/src/ocf_data_sampler/datasets/pvnet/preprocess.py
@@ -181,7 +181,7 @@ def apply_dropout_to_datasets(
     t0: np.datetime64,
     config: PVNetDataConfig,
 ) -> None:
-    """Apply dropout in-placeto the dictionary of input data sources around a given t0 time.
+    """Apply dropout in-place to the dictionary of input data sources around a given t0 time.
 
     Args:
         datasets_dict: Dictionary of the input data sources

From cfa6e8e4cae8d6ba0032756a1e5e033ee5bb2370 Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Fri, 7 Aug 2026 13:07:57 +0000
Subject: [PATCH 9/9] Improve test_apply_dropout_to_datasets

---
 tests/datasets/pvnet/test_preprocess.py | 95 ++++++++++++-------------
 1 file changed, 44 insertions(+), 51 deletions(-)

diff --git a/tests/datasets/pvnet/test_preprocess.py b/tests/datasets/pvnet/test_preprocess.py
index 74a002c9..245bbc7c 100644
--- a/tests/datasets/pvnet/test_preprocess.py
+++ b/tests/datasets/pvnet/test_preprocess.py
@@ -10,6 +10,27 @@
 )
 
 
+def _generation_da(generation_mw: np.ndarray, capacity_mwp: np.ndarray) -> xr.DataArray:
+    """Build a generation DataArray from 2D `(time_utc, location_id)` MW and capacity arrays."""
+    n_times, n_locations = generation_mw.shape
+    return xr.DataArray(
+        np.stack([generation_mw, capacity_mwp], axis=-1),
+        coords={
+            "time_utc": pd.date_range("2023-01-01", periods=n_times, freq="30min"),
+            "location_id": np.arange(1, n_locations + 1),
+            "gen_param": ["generation_mw", "capacity_mwp"],
+        },
+        dims=("time_utc", "location_id", "gen_param"),
+    )
+
+
+def _assert_dropout_applied(da: xr.DataArray, cutoff_time: np.datetime64):
+    """Assert that all values after the cutoff time are NaN, and all values before are not NaN.
+    """
+    assert not np.any(np.isnan(da.sel(time_utc=slice(None, cutoff_time))))
+    assert np.all(np.isnan(da.sel(time_utc=slice(cutoff_time + np.timedelta64(1, "s"), None))))
+
+
 def test_fill_nans_in_dataset_dicts(config_filename):
     """Test the fill_nans_in_arrays function from configuration"""
 
@@ -39,20 +60,6 @@ def test_fill_nans_in_dataset_dicts(config_filename):
     assert np.array_equal(datasets_dict["nwp"]["ukv"].values, np.array([-2.0, 3.0, -2.0]))
 
 
-def _generation_da(generation_mw: np.ndarray, capacity_mwp: np.ndarray) -> xr.DataArray:
-    """Build a generation DataArray from 2D `(time_utc, location_id)` MW and capacity arrays."""
-    n_times, n_locations = generation_mw.shape
-    return xr.DataArray(
-        np.stack([generation_mw, capacity_mwp], axis=-1),
-        coords={
-            "time_utc": pd.date_range("2023-01-01", periods=n_times, freq="30min"),
-            "location_id": np.arange(1, n_locations + 1),
-            "gen_param": ["generation_mw", "capacity_mwp"],
-        },
-        dims=("time_utc", "location_id", "gen_param"),
-    )
-
-
 def test_normalise_dataset_dicts_generation():
     """Generation is normalised to a capacity factor (generation_mw / capacity_mwp)."""
     capacity_mwp = np.array([[100.0, 0.0], [100.0, 40.0]])
@@ -84,55 +91,41 @@ def test_normalise_dataset_dicts_generation():
 def test_apply_dropout_to_datasets(pvnet_config_filename):
     config = load_yaml_configuration(pvnet_config_filename)
 
-    # Set dropout on the input window only - generation.target has no dropout config
     config.generation.input.dropout_timedeltas_minutes = [-30]
     config.generation.input.dropout_fraction = 1.0
-    config.satellite.dropout_timedeltas_minutes = []
-    config.satellite.dropout_fraction = 0
-
-    t0 = np.datetime64("2023-01-01 12:00")
-    times = np.array(
-        [
-            "2023-01-01T11:00",
-            "2023-01-01T11:30",
-            "2023-01-01T12:00",
-            "2023-01-01T12:30",
-        ],
-        dtype="datetime64[m]",
-    )
-    generation_mw = np.arange(4 * 2, dtype=float).reshape(4, 2)
-    capacity_mwp = np.ones((4, 2))
-    generation = xr.DataArray(
-        np.stack([generation_mw, capacity_mwp], axis=-1),
-        coords={
-            "time_utc": times,
-            "location_id": [1, 2],
-            "gen_param": ["generation_mw", "capacity_mwp"],
-        },
-        dims=("time_utc", "location_id", "gen_param"),
-    )
+
+    config.satellite.dropout_timedeltas_minutes = [-60]
+    config.satellite.dropout_fraction = 1.0
+
+    generation_input = _generation_da(generation_mw=np.ones((4, 2)),  capacity_mwp=np.ones((4, 2)))
+    generation_target = generation_input.copy(deep=True)
+
+    t0 = np.datetime64("2023-01-01 00:00")
+
+    times = pd.date_range(end=t0, periods=4, freq="30min").values
+
     sat = xr.DataArray(
         np.arange(4, dtype=float),
         coords={"time_utc": times},
         dims=("time_utc",),
     )
 
-    datasets_dict = {"generation_input": generation, "sat": sat}
+    datasets_dict = {
+        "generation_input": generation_input,
+        "generation_target": generation_target,
+        "sat": sat
+    }
 
     apply_dropout_to_datasets(datasets_dict, t0, config)
 
-    ds_gen = datasets_dict["generation_input"].sel(gen_param="generation_mw")
-    ds_cap = datasets_dict["generation_input"].sel(gen_param="capacity_mwp")
+    # Generation dropout with a -30 minute cutoff should blank everything from t0
+    _assert_dropout_applied(datasets_dict["generation_input"], t0 - np.timedelta64(30, "m"))
+
+    # No dropout is applied to generation.target, so it should have no NaNs
+    assert not np.any(np.isnan(datasets_dict["generation_target"]))
 
-    # Generation dropout with a -30 minute cutoff should blank everything from t0 onwards,
-    # including timesteps beyond t0.
-    assert not np.any(np.isnan(ds_gen.sel(time_utc=slice(None, "2023-01-01T11:30"))))
-    assert np.all(np.isnan(ds_gen.sel(time_utc=slice("2023-01-01T12:00", None))))
+    # Satellite dropout with a -60 minute cutoff should blank everything from t0 - 60 minutes
+    _assert_dropout_applied(datasets_dict["sat"], t0 - np.timedelta64(60, "m"))
 
-    # capacity_mwp is dropped out along with generation_mw, using the same cutoff.
-    assert not np.any(np.isnan(ds_cap.sel(time_utc=slice(None, "2023-01-01T11:30"))))
-    assert np.all(np.isnan(ds_cap.sel(time_utc=slice("2023-01-01T12:00", None))))
 
-    # Satellite dropout is disabled, so the helper should leave it untouched.
-    xr.testing.assert_equal(datasets_dict["sat"], sat)