Merge pull request #204 from ecmwf/feature/fix_local_regular_grid

Feature/fix local regular grid

polytope/datacube/backends/datacube.py

@@ -128,6 +128,7 @@ def get_indices(self, path: DatacubePath, axis, lower, upper, method=None):
         """
         path = self.fit_path(path)
         indexes = axis.find_indexes(path, self)
+
         idx_between = axis.find_indices_between(indexes, lower, upper, self, method)
 
         logging.debug(f"For axis {axis.name} between {lower} and {upper}, found indices {idx_between}")

polytope/datacube/transformations/datacube_mappers/datacube_mappers.py

@@ -15,6 +15,9 @@ def __init__(self, name, mapper_options):
         self.local_area = []
         if mapper_options.local is not None:
             self.local_area = mapper_options.local
+        self._axis_reversed = None
+        if mapper_options.axis_reversed is not None:
+            self._axis_reversed = mapper_options.axis_reversed
         self.old_axis = name
         self._final_transformation = self.generate_final_transformation()
         self._final_mapped_axes = self._final_transformation._mapped_axes
@@ -26,7 +29,9 @@ def generate_final_transformation(self):
         map_type = _type_to_datacube_mapper_lookup[self.grid_type]
         module = import_module("polytope.datacube.transformations.datacube_mappers.mapper_types." + self.grid_type)
         constructor = getattr(module, map_type)
-        transformation = deepcopy(constructor(self.old_axis, self.grid_axes, self.grid_resolution, self.local_area))
+        transformation = deepcopy(
+            constructor(self.old_axis, self.grid_axes, self.grid_resolution, self.local_area, self._axis_reversed)
+        )
         return transformation
 
     def blocked_axes(self):

polytope/datacube/transformations/datacube_mappers/mapper_types/healpix.py

@@ -5,7 +5,7 @@
 
 
 class HealpixGridMapper(DatacubeMapper):
-    def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
+    def __init__(self, base_axis, mapped_axes, resolution, local_area=[], axis_reversed=None):
         # TODO: if local area is not empty list, raise NotImplemented
         self._mapped_axes = mapped_axes
         self._base_axis = base_axis
@@ -14,6 +14,10 @@ def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
         self._first_axis_vals = self.first_axis_vals()
         self.compressed_grid_axes = [self._mapped_axes[1]]
         self.md5_hash = md5_hash.get(resolution, None)
+        if self._axis_reversed[mapped_axes[1]]:
+            raise NotImplementedError("Healpix grid with second axis in decreasing order is not supported")
+        if not self._axis_reversed[mapped_axes[0]]:
+            raise NotImplementedError("Healpix grid with first axis in increasing order is not supported")
 
     def first_axis_vals(self):
         rad2deg = 180 / math.pi

polytope/datacube/transformations/datacube_mappers/mapper_types/healpix_nested.py

@@ -5,7 +5,7 @@
 
 
 class NestedHealpixGridMapper(DatacubeMapper):
-    def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
+    def __init__(self, base_axis, mapped_axes, resolution, local_area=[], axis_reversed=None):
         # TODO: if local area is not empty list, raise NotImplemented
         self._mapped_axes = mapped_axes
         self._base_axis = base_axis
@@ -18,6 +18,10 @@ def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
         self.Npix = 12 * self.Nside * self.Nside
         self.Ncap = (self.Nside * (self.Nside - 1)) << 1
         self.md5_hash = md5_hash.get(resolution, None)
+        if self._axis_reversed[mapped_axes[1]]:
+            raise NotImplementedError("Healpix grid with second axis in decreasing order is not supported")
+        if not self._axis_reversed[mapped_axes[0]]:
+            raise NotImplementedError("Healpix grid with first axis in increasing order is not supported")
 
     def first_axis_vals(self):
         rad2deg = 180 / math.pi

polytope/datacube/transformations/datacube_mappers/mapper_types/local_regular.py

@@ -4,7 +4,7 @@
 
 
 class LocalRegularGridMapper(DatacubeMapper):
-    def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
+    def __init__(self, base_axis, mapped_axes, resolution, local_area=[], axis_reversed=None):
         # TODO: if local area is not empty list, raise NotImplemented
         self._mapped_axes = mapped_axes
         self._base_axis = base_axis
@@ -22,12 +22,25 @@ def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
             self.md5_hash = md5_hash.get(tuple(resolution), None)
         self._first_deg_increment = (local_area[1] - local_area[0]) / self.first_resolution
         self._second_deg_increment = (local_area[3] - local_area[2]) / self.second_resolution
-        self._axis_reversed = {mapped_axes[0]: True, mapped_axes[1]: False}
+        if axis_reversed is None:
+            self._axis_reversed = {mapped_axes[0]: False, mapped_axes[1]: False}
+        else:
+            assert set(axis_reversed.keys()) == set(mapped_axes)
+            self._axis_reversed = axis_reversed
         self._first_axis_vals = self.first_axis_vals()
         self.compressed_grid_axes = [self._mapped_axes[1]]
+        if self._axis_reversed[mapped_axes[1]]:
+            raise NotImplementedError("Local regular grid with second axis in decreasing order is not supported")
 
     def first_axis_vals(self):
-        first_ax_vals = [self._first_axis_max - i * self._first_deg_increment for i in range(self.first_resolution + 1)]
+        if self._axis_reversed[self._mapped_axes[0]]:
+            first_ax_vals = [
+                self._first_axis_max - i * self._first_deg_increment for i in range(self.first_resolution + 1)
+            ]
+        else:
+            first_ax_vals = [
+                self._first_axis_min + i * self._first_deg_increment for i in range(self.first_resolution + 1)
+            ]
         return first_ax_vals
 
     def map_first_axis(self, lower, upper):

polytope/datacube/transformations/datacube_mappers/mapper_types/octahedral.py

@@ -5,7 +5,7 @@
 
 
 class OctahedralGridMapper(DatacubeMapper):
-    def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
+    def __init__(self, base_axis, mapped_axes, resolution, local_area=[], axis_reversed=None):
         # TODO: if local area is not empty list, raise NotImplemented
         self._mapped_axes = mapped_axes
         self._base_axis = base_axis
@@ -14,6 +14,10 @@ def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
         self._first_idx_map = self.create_first_idx_map()
         self._second_axis_spacing = {}
         self._axis_reversed = {mapped_axes[0]: True, mapped_axes[1]: False}
+        if self._axis_reversed[mapped_axes[1]]:
+            raise NotImplementedError("Octahedral grid with second axis in decreasing order is not supported")
+        if not self._axis_reversed[mapped_axes[0]]:
+            raise NotImplementedError("Octahedral grid with first axis in increasing order is not supported")
         self.compressed_grid_axes = [self._mapped_axes[1]]
         self.md5_hash = md5_hash.get(resolution, None)
 

polytope/datacube/transformations/datacube_mappers/mapper_types/reduced_ll.py

@@ -4,7 +4,7 @@
 
 
 class ReducedLatLonMapper(DatacubeMapper):
-    def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
+    def __init__(self, base_axis, mapped_axes, resolution, local_area=[], axis_reversed=None):
         # TODO: if local area is not empty list, raise NotImplemented
         self._mapped_axes = mapped_axes
         self._base_axis = base_axis
@@ -13,6 +13,10 @@ def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
         self._first_axis_vals = self.first_axis_vals()
         self.compressed_grid_axes = [self._mapped_axes[1]]
         self.md5_hash = md5_hash.get(resolution, None)
+        if self._axis_reversed[mapped_axes[1]]:
+            raise NotImplementedError("Reduced lat-lon grid with second axis in decreasing order is not supported")
+        if self._axis_reversed[mapped_axes[0]]:
+            raise NotImplementedError("Reduced lat-lon grid with first axis in decreasing order is not supported")
 
     def first_axis_vals(self):
         resolution = 180 / (self._resolution - 1)

polytope/datacube/transformations/datacube_mappers/mapper_types/regular.py

@@ -4,19 +4,28 @@
 
 
 class RegularGridMapper(DatacubeMapper):
-    def __init__(self, base_axis, mapped_axes, resolution, local_area=[]):
+    def __init__(self, base_axis, mapped_axes, resolution, local_area=[], axis_reversed=None):
         # TODO: if local area is not empty list, raise NotImplemented
         self._mapped_axes = mapped_axes
         self._base_axis = base_axis
         self._resolution = resolution
         self.deg_increment = 90 / self._resolution
-        self._axis_reversed = {mapped_axes[0]: True, mapped_axes[1]: False}
+        if axis_reversed is None:
+            self._axis_reversed = {mapped_axes[0]: True, mapped_axes[1]: False}
+        else:
+            assert set(axis_reversed.keys()) == set(mapped_axes)
+            self._axis_reversed = axis_reversed
         self._first_axis_vals = self.first_axis_vals()
         self.compressed_grid_axes = [self._mapped_axes[1]]
         self.md5_hash = md5_hash.get(resolution, None)
+        if self._axis_reversed[mapped_axes[1]]:
+            raise NotImplementedError("Regular grid with second axis in decreasing order is not supported")
 
     def first_axis_vals(self):
-        first_ax_vals = [90 - i * self.deg_increment for i in range(2 * self._resolution)]
+        if self._axis_reversed[self._mapped_axes[0]]:
+            first_ax_vals = [90 - i * self.deg_increment for i in range(2 * self._resolution)]
+        else:
+            first_ax_vals = [-90 + i * self.deg_increment for i in range(2 * self._resolution)]
         return first_ax_vals
 
     def map_first_axis(self, lower, upper):

polytope/options.py

@@ -22,6 +22,7 @@ class MapperConfig(TransformationConfig):
     resolution: Union[int, List[int]] = 0
     axes: List[str] = [""]
     local: Optional[List[float]] = None
+    axis_reversed: Optional[Dict[str, bool]] = None
 
 
 class ReverseConfig(TransformationConfig):

tests/test_healpix_mapper.py

@@ -10,7 +10,7 @@
 from polytope.shapes import Box, Select
 
 
-class TestOctahedralGrid:
+class TestHealpixGrid:
     def setup_method(self, method):
         nexus_url = "https://get.ecmwf.int/test-data/polytope/test-data/healpix.grib"
         download_test_data(nexus_url, "healpix.grib")
@@ -69,7 +69,7 @@ def test_healpix_grid(self):
                 eccodes_lon = nearest_points[0][0]["lon"]
                 eccodes_result = nearest_points[0][0]["value"]
 
-                mapper = HealpixGridMapper("base", ["base", "base"], 32)
+                mapper = HealpixGridMapper("base", ["base1", "base2"], 32)
                 assert nearest_points[0][0]["index"] == mapper.unmap((lat,), (lon,))
                 assert eccodes_lat - tol <= lat
                 assert lat <= eccodes_lat + tol

tests/test_local_grid_cyclic.py

@@ -26,6 +26,7 @@ def setup_method(self, method):
                             "resolution": [80, 80],
                             "axes": ["latitude", "longitude"],
                             "local": [-40, 40, -20, 60],
+                            "axis_reversed": {"latitude": True, "longitude": False},
                         }
                     ],
                 },

tests/test_local_regular_grid.py

@@ -26,6 +26,7 @@ def setup_method(self, method):
                             "resolution": [80, 80],
                             "axes": ["latitude", "longitude"],
                             "local": [-40, 40, -20, 60],
+                            "axis_reversed": {"latitude": True, "longitude": False},
                         }
                     ],
                 },

tests/test_regular_grid.py

@@ -25,7 +25,13 @@ def setup_method(self, method):
                 {
                     "axis_name": "values",
                     "transformations": [
-                        {"name": "mapper", "type": "regular", "resolution": 30, "axes": ["latitude", "longitude"]}
+                        {
+                            "name": "mapper",
+                            "type": "regular",
+                            "resolution": 30,
+                            "axes": ["latitude", "longitude"],
+                            "axis_reversed": {"latitude": True, "longitude": False},
+                        }
                     ],
                 },
                 {"axis_name": "latitude", "transformations": [{"name": "reverse", "is_reverse": True}]},
@@ -108,7 +114,7 @@ def test_regular_grid(self):
             eccodes_value = nearest_points[121][0]["value"]
             eccodes_lats.append(eccodes_lat)
 
-            mapper = RegularGridMapper("base", ["base", "base"], 30)
+            mapper = RegularGridMapper("base", ["base1", "base2"], 30)
             assert nearest_points[121][0]["index"] == mapper.unmap((lat,), (lon,))
 
             assert eccodes_lat - tol <= lat