Merge pull request #640 from openclimatefix/issue/628-satellite-order

Issue/628 satellite order

Author: peterdudfield

nowcasting_dataset/data_sources/satellite/satellite_model.py

@@ -7,6 +7,14 @@
 
 logger = logging.getLogger(__name__)
 
+satellite_expected_dims_order = (
+    "example",
+    "time_index",
+    "channels_index",
+    "y_geostationary_index",
+    "x_geostationary_index",
+)
+
 
 class Satellite(DataSourceOutput):
     """Class to store satellite data as a xr.Dataset with some validation"""
@@ -34,6 +42,10 @@ def model_validation(cls, v):
         v.check_data_var_dim(v.x_geostationary, ("example", "x_geostationary_index"))
         v.check_data_var_dim(v.y_geostationary, ("example", "y_geostationary_index"))
 
+        # re-order dims
+        if v.data.dims != satellite_expected_dims_order:
+            v.__setitem__("data", v.data.transpose(*satellite_expected_dims_order))
+
         return v
 
 

tests/data_sources/satellite/test_satellite_model.py

@@ -6,7 +6,10 @@
 import pytest
 
 from nowcasting_dataset.data_sources.fake.batch import satellite_fake
-from nowcasting_dataset.data_sources.satellite.satellite_model import Satellite
+from nowcasting_dataset.data_sources.satellite.satellite_model import (
+    Satellite,
+    satellite_expected_dims_order,
+)
 
 
 def test_satellite_init():  # noqa: D103
@@ -15,6 +18,12 @@ def test_satellite_init():  # noqa: D103
     assert satellite.x_geostationary.dims == ("example", "x_geostationary_index")
 
 
+def test_satellite_dims_order():  # noqa: D103
+    satellite = satellite_fake()
+    satellite = Satellite.model_validation(satellite)
+    assert satellite.data.dims == satellite_expected_dims_order
+
+
 def test_satellite_validation():  # noqa: D103
     sat = satellite_fake()