feat(repositories): Enable working with eu stream

Author: devsjc

src/nwp_consumer/internal/repositories/raw_repositories/ecmwf_realtime.py

@@ -21,9 +21,10 @@
 time of 2024-10-25 00:00 and a target time of 2024-10-26 01:00 (step of 25 hours).
 
 The file contents is specific to the order agreed with the data provider.
-For the order that OCF has created, there are four distinct datasets.
-This is because OCF has ordered two separate regions and 17 variables,
-which are split across two datasets.
+OCF orders a single west-europe region alongside a separate india region,
+and the variables are split across multiple datasets within each file.
+Datasets that overlap the requested model region are cropped to it, which
+allows building a sub-region store (e.g. nl) from the larger west-europe order.
 
 Also, some of the data contains larger steps than we are interested in due
 to necessities in the order creation process.
@@ -91,6 +92,9 @@ def repository() -> entities.RawRepositoryMetadata:
                 "hres-ifs-nl": entities.Models.ECMWF_HRES_IFS_0P1DEGREE.with_region(
                     "nl",
                 ).with_max_step(84).with_delay_minutes(60 * 7),
+                "hres-ifs-west-europe": entities.Models.ECMWF_HRES_IFS_0P1DEGREE.with_region(
+                    "west-europe",
+                ).with_delay_minutes(60 * 7),
             },
         )
 
@@ -273,14 +277,22 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
         expected_lats = ECMWFRealTimeS3RawRepository.model().expected_coordinates.latitude
         expected_steps = ECMWFRealTimeS3RawRepository.model().expected_coordinates.step
 
+        north, west, south, east = \
+            ECMWFRealTimeS3RawRepository.model().expected_coordinates.nwse()
+
         for i, ds in enumerate(dss):
-            # ECMWF Realtime provides all regions in one set of datasets,
-            # so distinguish via their coordinates
+            # ECMWF Realtime provides multiple orders in one set of datasets
+            # (e.g. west-europe and india), so distinguish via their coordinates.
+            # Datasets that overlap the requested region are kept and cropped to it,
+            # which allows a sub-region store (e.g. nl) to be built from the
+            # larger west-europe order.
+            ds_lons = ds.coords["longitude"].values
+            ds_lats = ds.coords["latitude"].values
             step = np.timedelta64(ds.coords["step"].values, "h").astype(int) # type: ignore[arg-type]
             is_relevant_dataset_predicate: bool = (
                 (expected_lons is not None and expected_lats is not None)
-                and (expected_lons[0] <= max(ds.coords["longitude"].values) <= expected_lons[-1])
-                and (expected_lats[-1] <= max(ds.coords["latitude"].values) <= expected_lats[0])
+                and (min(ds_lons) <= expected_lons[-1] and max(ds_lons) >= expected_lons[0])
+                and (min(ds_lats) <= expected_lats[0] and max(ds_lats) >= expected_lats[-1])
                 and (expected_steps[0] <= step <= expected_steps[-1])
             )
             if not is_relevant_dataset_predicate:
@@ -311,6 +323,15 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
                     .sortby(variables=["step", "variable", "longitude"])
                     .sortby(variables="latitude", ascending=False)
                 )
+                # Crop to the requested model region, since the order may cover
+                # a larger extent (e.g. cropping nl from the west-europe order).
+                # Round coordinates to the store's 4 d.p. precision first, else
+                # floating-point grid-edge differences drop boundary points and
+                # misalign the write with chunk boundaries.
+                da = da.assign_coords(
+                    latitude=[float(f"{v:.4f}") for v in da["latitude"].values],
+                    longitude=[float(f"{v:.4f}") for v in da["longitude"].values],
+                ).sel(latitude=slice(north, south), longitude=slice(west, east))
 
             except Exception as e:
                 return Failure(

src/nwp_consumer/internal/repositories/raw_repositories/mo_datahub.py

@@ -168,6 +168,9 @@ def repository() -> entities.RawRepositoryMetadata:
                 "um-global-10km-nl": entities.Models.MO_UM_GLOBAL_10KM\
                     .with_region("nl")\
                     .with_extra_parameters([entities.Parameter.PRESSURE_MSL]),
+                "um-global-10km-west-europe": entities.Models.MO_UM_GLOBAL_10KM\
+                    .with_region("west-europe")\
+                    .with_extra_parameters([entities.Parameter.PRESSURE_MSL]),
                 "um-ukv-2km": entities.Models.MO_UM_UKV_2KM_LAEA,
             },
         )
@@ -403,6 +406,18 @@ def _convert_global(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
             da = da.sortby(variables=["step", "variable", "longitude"])
             da = da.sortby(variables="latitude", ascending=False)
 
+            # The order covers the full west-europe extent, so crop to the
+            # requested model region. Round coordinates to the store's 4 d.p.
+            # precision first, else floating-point grid-edge differences drop
+            # boundary points and misalign the write with chunk boundaries.
+            # This is a no-op when the model region matches the order bounds.
+            north, west, south, east = \
+                MetOfficeDatahubRawRepository.model().expected_coordinates.nwse()
+            da = da.assign_coords(
+                latitude=[float(f"{v:.4f}") for v in da["latitude"].values],
+                longitude=[float(f"{v:.4f}") for v in da["longitude"].values],
+            ).sel(latitude=slice(north, south), longitude=slice(west, east))
+
         except Exception as e:
             return Failure(
                 ValueError(f"Error processing DataArray for path '{path}'. Error context: {e}"),

src/nwp_consumer/internal/repositories/raw_repositories/test_ecmwf_realtime.py

@@ -8,7 +8,7 @@
 
 from returns.result import Failure, ResultE, Success
 
-from ...entities import NWPDimensionCoordinateMap, Parameter
+from ...entities import Models, NWPDimensionCoordinateMap, Parameter
 from .ecmwf_realtime import ECMWFRealTimeS3RawRepository
 
 if TYPE_CHECKING:
@@ -170,3 +170,35 @@ class TestCase:
     @patch.dict(os.environ, {"MODEL": "hres-ifs-india"}, clear=True)
     def test_convert_india(self) -> None:
         return self.test__convert()
+
+    def test_convert_crops_to_region(self) -> None:
+        """Test that _convert crops datasets to the requested model region.
+
+        The test grib contains a uk dataset (lon -12..3, lat 48..60); requesting
+        a smaller sub-region should crop the output down to it, mimicking cropping
+        e.g. nl from the larger west-europe order.
+        """
+        cropped_model = dataclasses.replace(
+            Models.ECMWF_HRES_IFS_0P1DEGREE,
+            name="hres-ifs_uk-sub",
+            expected_coordinates=Models.ECMWF_HRES_IFS_0P1DEGREE.expected_coordinates.crop(
+                north=55, west=-8, south=50, east=0,
+            ).unwrap(),
+        )
+
+        with patch.object(
+            ECMWFRealTimeS3RawRepository,
+            "model",
+            staticmethod(lambda: cropped_model),
+        ):
+            result = ECMWFRealTimeS3RawRepository._convert(
+                path=pathlib.Path(__file__).parent.absolute()
+                / "test_gribs"
+                / "test_ECMWFRealtime_HRES-IFS_10u_20241104T00_S60.grib",
+            )
+            self.assertIsInstance(result, Success, msg=f"{result!s}")
+            for da in result.unwrap():
+                self.assertGreaterEqual(float(da.latitude.min()), 50)
+                self.assertLessEqual(float(da.latitude.max()), 55)
+                self.assertGreaterEqual(float(da.longitude.min()), -8)
+                self.assertLessEqual(float(da.longitude.max()), 0)

src/nwp_consumer/internal/repositories/raw_repositories/test_mo_datahub.py

@@ -7,7 +7,7 @@
 
 from returns.result import Failure, ResultE, Success
 
-from ...entities import NWPDimensionCoordinateMap, Parameter
+from ...entities import Models, NWPDimensionCoordinateMap, Parameter
 from .mo_datahub import MetOfficeDatahubRawRepository
 
 
@@ -95,6 +95,47 @@ class TestCase:
 
 
 
+    def test_convert_crops_to_region(self) -> None:
+        """Test that _convert_global crops the data to the requested model region.
+
+        The order provides a larger extent (here, the india test grib), which
+        should be cropped down to a smaller requested region.
+        """
+        cropped_model = dataclasses.replace(
+            Models.MO_UM_GLOBAL_10KM,
+            name="um-global_india-sub",
+            expected_coordinates=Models.MO_UM_GLOBAL_10KM.expected_coordinates.crop(
+                north=20, west=70, south=10, east=90,
+            ).unwrap(),
+        )
+
+        with patch.object(
+            MetOfficeDatahubRawRepository,
+            "model",
+            staticmethod(lambda: cropped_model),
+        ):
+            result = MetOfficeDatahubRawRepository._convert_global(
+                path=pathlib.Path(__file__).parent.absolute()
+                / "test_gribs"
+                / "test_MODatahub_UM-Global_t2m_20241120T00_S00.grib",
+            )
+            self.assertIsInstance(result, Success, msg=f"{result!s}")
+            da = result.unwrap()[0]
+            self.assertGreaterEqual(float(da.latitude.min()), 10)
+            self.assertLessEqual(float(da.latitude.max()), 20)
+            self.assertGreaterEqual(float(da.longitude.min()), 70)
+            self.assertLessEqual(float(da.longitude.max()), 90)
+            # Output must be a contiguous subset of the cropped region store
+            region_result = NWPDimensionCoordinateMap.from_xarray(da).bind(
+                dataclasses.replace(
+                    cropped_model.expected_coordinates,
+                    init_time=[dt.datetime(2024, 11, 20, 0, tzinfo=dt.UTC)],
+                    variable=[Parameter.TEMPERATURE_SL],
+                    step=[0],
+                ).determine_region,
+            )
+            self.assertIsInstance(region_result, Success, msg=f"{region_result!s}")
+
     @patch.dict(os.environ, {"MODEL": "um-ukv-2km"}, clear=True)
     def test_convert_ukv(self) -> None:
         @dataclasses.dataclass