Add fine coordinates to the model for easier inference handling

fme/core/coordinates.py

@@ -619,7 +619,7 @@ def __eq__(self, other) -> bool:
     def __repr__(self) -> str:
         return f"LatLonCoordinates(\n    lat={self.lat},\n    lon={self.lon}\n"
 
-    def to(self, device: str) -> "LatLonCoordinates":
+    def to(self, device: str | torch.device) -> "LatLonCoordinates":
         return LatLonCoordinates(
             lon=self.lon.to(device),
             lat=self.lat.to(device),

fme/downscaling/data/__init__.py

@@ -12,7 +12,12 @@
     PairedBatchItem,
     PairedGriddedData,
 )
-from .static import StaticInput, StaticInputs, load_static_inputs
+from .static import (
+    StaticInput,
+    StaticInputs,
+    load_fine_coords_from_path,
+    load_static_inputs,
+)
 from .utils import (
     BatchedLatLonCoordinates,
     ClosedInterval,

fme/downscaling/data/config.py

@@ -23,7 +23,11 @@
     PairedBatchData,
     PairedGriddedData,
 )
-from fme.downscaling.data.utils import ClosedInterval, adjust_fine_coord_range
+from fme.downscaling.data.utils import (
+    ClosedInterval,
+    adjust_fine_coord_range,
+    get_latlon_coords_from_properties,
+)
 from fme.downscaling.requirements import DataRequirements
 
 
@@ -529,6 +533,7 @@ def build(
             dims=example.fine.latlon_coordinates.dims,
             variable_metadata=variable_metadata,
             all_times=all_times,
+            fine_coords=get_latlon_coords_from_properties(properties_fine),
         )
 
     def _get_sampler(

fme/downscaling/data/datasets.py

@@ -138,11 +138,13 @@ def __init__(
                 f"expected lon_min < {self.lon_interval.start + 360.0}"
             )
 
-        self._lats_slice = self.lat_interval.slice_of(self._orig_coords.lat)
-        self._lons_slice = self.lon_interval.slice_of(self._orig_coords.lon)
+        # Used to subset the data in __getitem__
+        self._lats_slice = self.lat_interval.slice_from(self._orig_coords.lat)
+        self._lons_slice = self.lon_interval.slice_from(self._orig_coords.lon)
+
         self._latlon_coordinates = LatLonCoordinates(
-            lat=self._orig_coords.lat[self._lats_slice],
-            lon=self._orig_coords.lon[self._lons_slice],
+            lat=self.lat_interval.subset_of(self._orig_coords.lat),
+            lon=self.lon_interval.subset_of(self._orig_coords.lon),
         )
         self._area_weights = self._latlon_coordinates.area_weights
 
@@ -337,6 +339,7 @@ class PairedGriddedData:
     dims: list[str]
     variable_metadata: Mapping[str, VariableMetadata]
     all_times: xr.CFTimeIndex
+    fine_coords: LatLonCoordinates
 
     @property
     def loader(self) -> DataLoader[PairedBatchItem]: