Pad Satellite data with zeros (#653)

* Pad data with zeros

* Change padding

* Fix checks and add test

* Fix path

Author: jacobbieker

nowcasting_dataset/data_sources/satellite/satellite_data_source.py

@@ -162,67 +162,47 @@ def get_spatial_region_of_interest(
         max_x_and_y_index_height = x_and_y_index_at_center + half_image_size_pixels_height
 
         # Check whether the requested region of interest steps outside of the available data:
-        suggested_reduction_of_image_size_pixels_width = (
-            max(
-                (-min_x_and_y_index_width.min() if (min_x_and_y_index_width < 0).any() else 0),
-                (min_x_and_y_index_width.x_index_at_center - len(data_array.x_geostationary)),
-            )
-            * 2
+        # Need to know how much to pad the outputs, so can do that here
+        min_width_padding = max(-min_x_and_y_index_width.x_index_at_center, 0)
+        max_width_padding = max(
+            max_x_and_y_index_width.x_index_at_center - len(data_array.x_geostationary), 0
         )
-        suggested_reduction_of_image_size_pixels_height = (
-            max(
-                (-max_x_and_y_index_height.min() if (max_x_and_y_index_height < 0).any() else 0),
-                (max_x_and_y_index_height.y_index_at_center - len(data_array.y_geostationary)),
-            )
-            * 2
+        min_height_padding = max(-min_x_and_y_index_height.y_index_at_center, 0)
+        max_height_padding = max(
+            max_x_and_y_index_height.y_index_at_center - len(data_array.y_geostationary), 0
         )
-        # If the requested region does step outside the available data then raise an exception
-        # with a helpful message:
-        if (
-            suggested_reduction_of_image_size_pixels_width > 0
-            or suggested_reduction_of_image_size_pixels_height > 0
-        ):
-            new_suggested_image_size_pixels_width = (
-                self._rectangle.size_pixels_width - suggested_reduction_of_image_size_pixels_width
-            )
-            new_suggested_image_size_pixels_height = (
-                self._rectangle.size_pixels_height - suggested_reduction_of_image_size_pixels_height
-            )
-            raise RuntimeError(
-                "Requested region of interest of satellite data steps outside of the available"
-                " geographical extent of the Zarr data.  The requested region of interest extends"
-                f" from pixel indicies"
-                f" x={min_x_and_y_index_width.x_index_at_center} to"
-                f" x={min_x_and_y_index_width.x_index_at_center},"
-                f" y={min_x_and_y_index_width.y_index_at_center} to"
-                f" y={min_x_and_y_index_width.y_index_at_center}."
-                f" In the Zarr data, len(x)={len(data_array.x_geostationary)},"
-                f" len(y)={len(data_array.y_geostationary)}."
-                f" Try reducing image_size_pixels_height from {self._rectangle.size_pixels_height}"
-                f" to {new_suggested_image_size_pixels_height} pixels."
-                f" Try reducing image_size_pixels_width from {self._rectangle.size_pixels_width}"
-                f" to {new_suggested_image_size_pixels_width} pixels."
-                f" {self.history_length=}; {self.forecast_length=}; {x_center_osgb=};"
-                f" {y_center_osgb=}; {x_center_geostationary=}; {y_center_geostationary=};"
-                f" {min(data_array.x_geostationary.values)=};"
-                f" {max(data_array.x_geostationary.values)=};"
-                f" {min(data_array.y_geostationary.values)=};"
-                f" {max(data_array.y_geostationary.values)=};\n"
-                f" {data_array=}"
-            )
 
         # Select the geographical region of interest.
         # Note that isel is *exclusive* of the end of the slice.
         # e.g. isel(x=slice(0, 3)) will return the first, second, and third values.
         data_array = data_array.isel(
             x_geostationary=slice(
-                min_x_and_y_index_width.x_index_at_center, max_x_and_y_index_width.x_index_at_center
+                max(min_x_and_y_index_width.x_index_at_center, 0),
+                min(max_x_and_y_index_width.x_index_at_center, len(data_array.x_geostationary)),
             ),
             y_geostationary=slice(
-                min_x_and_y_index_height.y_index_at_center,
-                max_x_and_y_index_height.y_index_at_center,
+                max(min_x_and_y_index_height.y_index_at_center, 0),
+                min(max_x_and_y_index_height.y_index_at_center, len(data_array.y_geostationary)),
             ),
         )
+
+        # isel is wrong if padding before, so add padding after to be the correct size
+        # Get the difference for each direction and use that
+        if (
+            min_height_padding > 0
+            or min_width_padding > 0
+            or max_height_padding > 0
+            or max_width_padding > 0
+        ):
+            data_array = data_array.pad(
+                pad_width={
+                    "x_geostationary": (min_width_padding, max_width_padding),
+                    "y_geostationary": (min_height_padding, max_height_padding),
+                },
+                mode="constant",
+                constant_values=0,
+            )
+
         return data_array
 
     def get_example(self, location: SpaceTimeLocation) -> xr.Dataset:

tests/data_sources/satellite/test_satellite_data_source.py

@@ -55,6 +55,27 @@ def _test_get_example(
     assert len(sat_data.x_geostationary.shape) == 1
 
 
+def test_padding(sat_filename):  # noqa: D103
+    data_source = SatelliteDataSource(
+        image_size_pixels_height=1024,
+        image_size_pixels_width=1024,
+        zarr_path=sat_filename,
+        history_minutes=0,
+        forecast_minutes=15,
+        channels=("IR_016",),
+        meters_per_pixel=6000,
+    )
+    data_source.open()
+    t0_dt = pd.Timestamp("2020-04-01T13:00")
+    sat_data = data_source.get_example(
+        SpaceTimeLocation(t0_datetime_utc=t0_dt, x_center_osgb=0, y_center_osgb=0)
+    )
+
+    assert len(sat_data.x_geostationary) == 1024
+    assert len(sat_data.y_geostationary) == 1024
+    assert len(sat_data.x_geostationary.shape) == 1
+
+
 @pytest.mark.parametrize(
     "x_center_osgb, y_center_osgb, left_geostationary, right_geostationary,"
     " top_geostationary, bottom_geostationary",