Downscaling model handles static input instead of DataLoader

fme/downscaling/data/datasets.py

@@ -441,6 +441,16 @@ def __post_init__(self):
     def horizontal_shape(self) -> tuple[int, int]:
         return self._horizontal_shape
 
+    @property
+    def lat_interval(self) -> ClosedInterval:
+        lat = self.latlon_coordinates.lat[0]  # all batch members identical; use first
+        return ClosedInterval(lat.min().item(), lat.max().item())
+
+    @property
+    def lon_interval(self) -> ClosedInterval:
+        lon = self.latlon_coordinates.lon[0]  # all batch members identical; use first
+        return ClosedInterval(lon.min().item(), lon.max().item())
+
     @classmethod
     def from_sequence(
         cls,

fme/downscaling/data/test_utils.py

@@ -49,6 +49,23 @@ def test_adjust_fine_coord_range(downscale_factor, lat_range):
     assert len(subsel_fine_lat) / len(subsel_coarse_lat) == downscale_factor
 
 
+def test_adjust_fine_coord_range_raises_near_domain_boundary():
+    downscale_factor = 4  # n_half_fine = 2
+    coarse_edges = torch.linspace(0, 6, 7)
+    coarse_lat = _fine_midpoints(coarse_edges, 1)
+    fine_lat = _fine_midpoints(coarse_edges, downscale_factor)
+    # Drop the first fine point so only 1 fine point exists below coarse_min=0.5,
+    # but n_half_fine=2 are required — simulating a grid truncated at the domain edge.
+    fine_lat_truncated = fine_lat[1:]
+    with pytest.raises(ValueError):
+        adjust_fine_coord_range(
+            ClosedInterval(0, 4),
+            full_coarse_coord=coarse_lat,
+            full_fine_coord=fine_lat_truncated,
+            downscale_factor=downscale_factor,
+        )
+
+
 @pytest.mark.parametrize(
     "input_slice,expected",
     [

fme/downscaling/data/utils.py

@@ -116,6 +116,13 @@ def adjust_fine_coord_range(
 
     If downscale factor is not provided, it is assumed that the coarse and fine
     coordinate tensors correspond to the same region bounds.
+
+    Raises:
+        ValueError: If coord_range is too close to the boundary of full_fine_coord
+            such that fewer than downscale_factor // 2 fine points exist beyond the
+            outermost selected coarse point on either side. For global latitude grids,
+            this is avoided by restricting coord_range to within ±88° (i.e. away from
+            the poles).
     """
     if downscale_factor is None:
         if full_fine_coord.shape[0] % full_coarse_coord.shape[0] != 0:
@@ -132,6 +139,19 @@ def adjust_fine_coord_range(
     n_half_fine = downscale_factor // 2
     coarse_min = full_coarse_coord[full_coarse_coord >= coord_range.start][0]
     coarse_max = full_coarse_coord[full_coarse_coord <= coord_range.stop][-1]
+
+    n_fine_below = int((full_fine_coord < coarse_min).sum())
+    n_fine_above = int((full_fine_coord > coarse_max).sum())
+    if n_fine_below < n_half_fine or n_fine_above < n_half_fine:
+        raise ValueError(
+            f"coord_range {coord_range} is too close to the boundary of "
+            f"full_fine_coord [{full_fine_coord.min():.2f}, "
+            f"{full_fine_coord.max():.2f}]. Need at least {n_half_fine} fine "
+            f"point(s) beyond each coarse boundary; got {n_fine_below} below "
+            f"and {n_fine_above} above. Restrict the coordinate range away from "
+            f"the domain edges."
+        )
+
     fine_min = full_fine_coord[full_fine_coord < coarse_min][-n_half_fine]
     fine_max = full_fine_coord[full_fine_coord > coarse_max][n_half_fine - 1]
 

fme/downscaling/models.py

@@ -15,8 +15,10 @@
 from fme.core.typing_ import TensorDict, TensorMapping
 from fme.downscaling.data import (
     BatchData,
+    ClosedInterval,
     PairedBatchData,
     StaticInputs,
+    adjust_fine_coord_range,
     load_static_inputs,
 )
 from fme.downscaling.metrics_and_maths import filter_tensor_mapping, interpolate
@@ -303,12 +305,33 @@ def __init__(
         self.out_packer = Packer(config.out_names)
         self.config = config
         self._channel_axis = -3
-        self.static_inputs = static_inputs
+        self.static_inputs = (
+            static_inputs.to_device() if static_inputs is not None else None
+        )
 
     @property
     def modules(self) -> torch.nn.ModuleList:
         return torch.nn.ModuleList([self.module])
 
+    def _subset_static_inputs(
+        self,
+        lat_interval: ClosedInterval,
+        lon_interval: ClosedInterval,
+    ) -> StaticInputs | None:
+        """Subset self.static_inputs to the given fine lat/lon interval.
+
+        Returns None if use_fine_topography is False.
+        Raises ValueError if use_fine_topography is True but self.static_inputs is None.
+        """
+        if not self.config.use_fine_topography:
+            return None
+        if self.static_inputs is None:
+            raise ValueError(
+                "Static inputs must be provided for each batch when use of fine "
+                "static inputs is enabled."
+            )
+        return self.static_inputs.subset_latlon(lat_interval, lon_interval)
+
     @property
     def fine_shape(self) -> tuple[int, int]:
         return self._get_fine_shape(self.coarse_shape)
@@ -338,6 +361,12 @@ def _get_input_from_coarse(
                     "static inputs is enabled."
                 )
             else:
+                expected_shape = interpolated.shape[-2:]
+                if static_inputs.shape != expected_shape:
+                    raise ValueError(
+                        f"Subsetted static input shape {static_inputs.shape} does not "
+                        f"match expected fine spatial shape {expected_shape}."
+                    )
                 n_batches = normalized.shape[0]
                 # Join normalized static inputs to input (see dataset for details)
                 for field in static_inputs.fields:
@@ -354,12 +383,17 @@ def _get_input_from_coarse(
     def train_on_batch(
         self,
         batch: PairedBatchData,
-        static_inputs: StaticInputs | None,
+        static_inputs: StaticInputs | None,  # TODO: remove in follow-on PR
         optimizer: Optimization | NullOptimization,
     ) -> ModelOutputs:
         """Performs a denoising training step on a batch of data."""
+        # Ignore the passed static_inputs; subset self.static_inputs using fine batch
+        # coordinates. The caller-provided value is kept for signature compatibility.
+        _static_inputs = self._subset_static_inputs(
+            batch.fine.lat_interval, batch.fine.lon_interval
+        )
         coarse, fine = batch.coarse.data, batch.fine.data
-        inputs_norm = self._get_input_from_coarse(coarse, static_inputs)
+        inputs_norm = self._get_input_from_coarse(coarse, _static_inputs)
         targets_norm = self.out_packer.pack(
             self.normalizer.fine.normalize(dict(fine)), axis=self._channel_axis
         )
@@ -415,9 +449,11 @@ def train_on_batch(
     def generate(
         self,
         coarse_data: TensorMapping,
-        static_inputs: StaticInputs | None,
+        static_inputs: StaticInputs | None,  # TODO: remove in follow-on PR
         n_samples: int = 1,
     ) -> tuple[TensorDict, torch.Tensor, list[torch.Tensor]]:
+        # static_inputs receives an internally-subsetted value from the calling method;
+        # external callers should use generate_on_batch / generate_on_batch_no_target.
         inputs_ = self._get_input_from_coarse(coarse_data, static_inputs)
         # expand samples and fold to
         # [batch * n_samples, output_channels, height, width]
@@ -467,22 +503,54 @@ def generate(
     def generate_on_batch_no_target(
         self,
         batch: BatchData,
-        static_inputs: StaticInputs | None,
+        static_inputs: StaticInputs | None,  # TODO: remove in follow-on PR
         n_samples: int = 1,
     ) -> TensorDict:
-        generated, _, _ = self.generate(batch.data, static_inputs, n_samples)
+        # Ignore the passed static_inputs; derive the fine lat/lon interval from coarse
+        # batch coordinates via adjust_fine_coord_range, then subset self.static_inputs.
+        if self.config.use_fine_topography:
+            if self.static_inputs is None:
+                raise ValueError(
+                    "Static inputs must be provided for each batch when use of fine "
+                    "static inputs is enabled."
+                )
+            coarse_lat = batch.latlon_coordinates.lat[0]
+            coarse_lon = batch.latlon_coordinates.lon[0]
+            fine_lat_interval = adjust_fine_coord_range(
+                batch.lat_interval,
+                full_coarse_coord=coarse_lat,
+                full_fine_coord=self.static_inputs.coords.lat,
+                downscale_factor=self.downscale_factor,
+            )
+            fine_lon_interval = adjust_fine_coord_range(
+                batch.lon_interval,
+                full_coarse_coord=coarse_lon,
+                full_fine_coord=self.static_inputs.coords.lon,
+                downscale_factor=self.downscale_factor,
+            )
+            _static_inputs = self.static_inputs.subset_latlon(
+                fine_lat_interval, fine_lon_interval
+            )
+        else:
+            _static_inputs = None
+        generated, _, _ = self.generate(batch.data, _static_inputs, n_samples)
         return generated
 
     @torch.no_grad()
     def generate_on_batch(
         self,
         batch: PairedBatchData,
-        static_inputs: StaticInputs | None,
+        static_inputs: StaticInputs | None,  # TODO: remove in follow-on PR
         n_samples: int = 1,
     ) -> ModelOutputs:
+        # Ignore the passed static_inputs; subset self.static_inputs using fine batch
+        # coordinates. The caller-provided value is kept for signature compatibility.
+        _static_inputs = self._subset_static_inputs(
+            batch.fine.lat_interval, batch.fine.lon_interval
+        )
         coarse, fine = batch.coarse.data, batch.fine.data
         generated, generated_norm, latent_steps = self.generate(
-            coarse, static_inputs, n_samples
+            coarse, _static_inputs, n_samples
         )
 
         targets_norm = self.out_packer.pack(