Merge pull request #70 from ecmwf/feature/no_cyclic_offsets

Surrounding method for surrounding points in Select

polytope/datacube/backends/datacube.py

@@ -85,7 +85,7 @@ def fit_path(self, path):
                 path.pop(key)
         return path
 
-    def get_indices(self, path: DatacubePath, axis, lower, upper):
+    def get_indices(self, path: DatacubePath, axis, lower, upper, method=None):
         """
         Given a path to a subset of the datacube, return the discrete indexes which exist between
         two non-discrete values (lower, upper) for a particular axis (given by label)
@@ -101,30 +101,34 @@ def get_indices(self, path: DatacubePath, axis, lower, upper):
         for r in original_search_ranges:
             offset = axis.offset(r)
             search_ranges_offset.append(offset)
-        idx_between = self._look_up_datacube(search_ranges, search_ranges_offset, indexes, axis)
+        idx_between = self._look_up_datacube(search_ranges, search_ranges_offset, indexes, axis, method)
         # Remove duplicates even if difference of the order of the axis tolerance
         if offset is not None:
             # Note that we can only do unique if not dealing with time values
             idx_between = unique(idx_between)
         return idx_between
 
-    def _look_up_datacube(self, search_ranges, search_ranges_offset, indexes, axis):
+    def _look_up_datacube(self, search_ranges, search_ranges_offset, indexes, axis, method):
         idx_between = []
         for i in range(len(search_ranges)):
             r = search_ranges[i]
             offset = search_ranges_offset[i]
             low = r[0]
             up = r[1]
-            indexes_between = axis.find_indices_between([indexes], low, up, self)
+            indexes_between = axis.find_indices_between([indexes], low, up, self, method)
             # Now the indexes_between are values on the cyclic range so need to remap them to their original
             # values before returning them
             for j in range(len(indexes_between)):
-                for k in range(len(indexes_between[j])):
-                    if offset is None:
-                        indexes_between[j][k] = indexes_between[j][k]
-                    else:
-                        indexes_between[j][k] = round(indexes_between[j][k] + offset, int(-math.log10(axis.tol)))
-                    idx_between.append(indexes_between[j][k])
+                # if we have a special indexes between range that needs additional offset, treat it here
+                if len(indexes_between[j]) == 0:
+                    idx_between = idx_between
+                else:
+                    for k in range(len(indexes_between[j])):
+                        if offset is None:
+                            indexes_between[j][k] = indexes_between[j][k]
+                        else:
+                            indexes_between[j][k] = round(indexes_between[j][k] + offset, int(-math.log10(axis.tol)))
+                        idx_between.append(indexes_between[j][k])
         return idx_between
 
     def get_mapper(self, axis):

polytope/datacube/backends/mock.py

@@ -41,7 +41,7 @@ def get(self, requests: IndexTree):
     def get_mapper(self, axis):
         return self.mappers[axis]
 
-    def get_indices(self, path: DatacubePath, axis, lower, upper):
+    def get_indices(self, path: DatacubePath, axis, lower, upper, method=None):
         if lower == upper == math.ceil(lower):
             if lower >= 0:
                 return [int(lower)]

polytope/datacube/datacube_axis.py

@@ -135,6 +135,37 @@ def unmap_to_datacube(path, unmapped_path):
             (path, unmapped_path) = old_unmap_to_datacube(path, unmapped_path)
             return (path, unmapped_path)
 
+        old_find_indices_between = cls.find_indices_between
+
+        def find_indices_between(index_ranges, low, up, datacube, method=None):
+            update_range()
+            indexes_between_ranges = []
+
+            if method != "surrounding":
+                return old_find_indices_between(index_ranges, low, up, datacube, method)
+            else:
+                for indexes in index_ranges:
+                    if cls.name in datacube.complete_axes:
+                        start = indexes.searchsorted(low, "left")
+                        end = indexes.searchsorted(up, "right")
+                    else:
+                        start = indexes.index(low)
+                        end = indexes.index(up)
+                    if start - 1 < 0:
+                        index_val_found = indexes[-1:][0]
+                        indexes_between_ranges.append([index_val_found])
+                    if end + 1 > len(indexes):
+                        index_val_found = indexes[:2][0]
+                        indexes_between_ranges.append([index_val_found])
+                    start = max(start - 1, 0)
+                    end = min(end + 1, len(indexes))
+                    if cls.name in datacube.complete_axes:
+                        indexes_between = indexes[start:end].to_list()
+                    else:
+                        indexes_between = indexes[start:end]
+                    indexes_between_ranges.append(indexes_between)
+                return indexes_between_ranges
+
         def offset(range):
             # We first unpad the range by the axis tolerance to make sure that
             # we find the wanted range of the cyclic axis since we padded by the axis tolerance before.
@@ -150,6 +181,7 @@ def offset(range):
         cls.find_indexes = find_indexes
         cls.unmap_to_datacube = unmap_to_datacube
         cls.unmap_total_path_to_datacube = unmap_total_path_to_datacube
+        cls.find_indices_between = find_indices_between
 
     return cls
 
@@ -235,15 +267,24 @@ def unmap_total_path_to_datacube(path, unmapped_path):
         def remap_to_requested(path, unmapped_path):
             return (path, unmapped_path)
 
-        def find_indices_between(index_ranges, low, up, datacube):
+        def find_indices_between(index_ranges, low, up, datacube, method=None):
+            # TODO: add method for snappping
             indexes_between_ranges = []
             for transform in cls.transformations:
                 if isinstance(transform, DatacubeMapper):
                     transformation = transform
                     if cls.name in transformation._mapped_axes():
                         for idxs in index_ranges:
-                            indexes_between = [i for i in idxs if low <= i <= up]
-                            indexes_between_ranges.append(indexes_between)
+                            if method == "surrounding":
+                                start = idxs.index(low)
+                                end = idxs.index(up)
+                                start = max(start - 1, 0)
+                                end = min(end + 1, len(idxs))
+                                indexes_between = idxs[start:end]
+                                indexes_between_ranges.append(indexes_between)
+                            else:
+                                indexes_between = [i for i in idxs if low <= i <= up]
+                                indexes_between_ranges.append(indexes_between)
             return indexes_between_ranges
 
         old_remap = cls.remap
@@ -307,15 +348,24 @@ def unmap_to_datacube(path, unmapped_path):
         def remap_to_requested(path, unmapped_path):
             return (path, unmapped_path)
 
-        def find_indices_between(index_ranges, low, up, datacube):
+        def find_indices_between(index_ranges, low, up, datacube, method=None):
+            # TODO: add method for snappping
             indexes_between_ranges = []
             for transform in cls.transformations:
                 if isinstance(transform, DatacubeAxisMerger):
                     transformation = transform
                     if cls.name in transformation._mapped_axes():
                         for indexes in index_ranges:
-                            indexes_between = [i for i in indexes if low <= i <= up]
-                            indexes_between_ranges.append(indexes_between)
+                            if method == "surrounding":
+                                start = indexes.index(low)
+                                end = indexes.index(up)
+                                start = max(start - 1, 0)
+                                end = min(end + 1, len(indexes))
+                                indexes_between = indexes[start:end]
+                                indexes_between_ranges.append(indexes_between)
+                            else:
+                                indexes_between = [i for i in indexes if low <= i <= up]
+                                indexes_between_ranges.append(indexes_between)
             return indexes_between_ranges
 
         def remap(range):
@@ -352,15 +402,42 @@ def unmap_to_datacube(path, unmapped_path):
         def remap_to_requested(path, unmapped_path):
             return (path, unmapped_path)
 
-        def find_indices_between(index_ranges, low, up, datacube):
+        def find_indices_between(index_ranges, low, up, datacube, method=None):
+            # TODO: add method for snappping
             indexes_between_ranges = []
             for transform in cls.transformations:
                 if isinstance(transform, DatacubeAxisReverse):
                     transformation = transform
                     if cls.name == transformation.name:
                         for indexes in index_ranges:
-                            indexes_between = [i for i in indexes if low <= i <= up]
-                            indexes_between_ranges.append(indexes_between)
+                            if cls.name in datacube.complete_axes:
+                                # Find the range of indexes between lower and upper
+                                # https://pandas.pydata.org/docs/reference/api/pandas.Index.searchsorted.html
+                                # Assumes the indexes are already sorted (could sort to be sure) and monotonically
+                                # increasing
+                                if method == "surrounding":
+                                    start = indexes.searchsorted(low, "left")
+                                    end = indexes.searchsorted(up, "right")
+                                    start = max(start - 1, 0)
+                                    end = min(end + 1, len(indexes))
+                                    indexes_between = indexes[start:end].to_list()
+                                    indexes_between_ranges.append(indexes_between)
+                                else:
+                                    start = indexes.searchsorted(low, "left")
+                                    end = indexes.searchsorted(up, "right")
+                                    indexes_between = indexes[start:end].to_list()
+                                    indexes_between_ranges.append(indexes_between)
+                            else:
+                                if method == "surrounding":
+                                    start = indexes.index(low)
+                                    end = indexes.index(up)
+                                    start = max(start - 1, 0)
+                                    end = min(end + 1, len(indexes))
+                                    indexes_between = indexes[start:end]
+                                    indexes_between_ranges.append(indexes_between)
+                                else:
+                                    indexes_between = [i for i in indexes if low <= i <= up]
+                                    indexes_between_ranges.append(indexes_between)
             return indexes_between_ranges
 
         def remap(range):
@@ -419,15 +496,24 @@ def unmap_to_datacube(path, unmapped_path):
         def remap_to_requested(path, unmapped_path):
             return (path, unmapped_path)
 
-        def find_indices_between(index_ranges, low, up, datacube):
+        def find_indices_between(index_ranges, low, up, datacube, method=None):
+            # TODO: add method for snappping
             indexes_between_ranges = []
             for transform in cls.transformations:
                 if isinstance(transform, DatacubeAxisTypeChange):
                     transformation = transform
                     if cls.name == transformation.name:
                         for indexes in index_ranges:
-                            indexes_between = [i for i in indexes if low <= i <= up]
-                            indexes_between_ranges.append(indexes_between)
+                            if method == "surrounding":
+                                start = indexes.index(low)
+                                end = indexes.index(up)
+                                start = max(start - 1, 0)
+                                end = min(end + 1, len(indexes))
+                                indexes_between = indexes[start:end]
+                                indexes_between_ranges.append(indexes_between)
+                            else:
+                                indexes_between = [i for i in indexes if low <= i <= up]
+                                indexes_between_ranges.append(indexes_between)
             return indexes_between_ranges
 
         def remap(range):
@@ -500,20 +586,37 @@ def unmap_total_path_to_datacube(self, path, unmapped_path):
     def remap_to_requeest(path, unmapped_path):
         return (path, unmapped_path)
 
-    def find_indices_between(self, index_ranges, low, up, datacube):
+    def find_indices_between(self, index_ranges, low, up, datacube, method=None):
+        # TODO: add method for snappping
         indexes_between_ranges = []
         for indexes in index_ranges:
             if self.name in datacube.complete_axes:
                 # Find the range of indexes between lower and upper
                 # https://pandas.pydata.org/docs/reference/api/pandas.Index.searchsorted.html
                 # Assumes the indexes are already sorted (could sort to be sure) and monotonically increasing
-                start = indexes.searchsorted(low, "left")
-                end = indexes.searchsorted(up, "right")
-                indexes_between = indexes[start:end].to_list()
-                indexes_between_ranges.append(indexes_between)
+                if method == "surrounding":
+                    start = indexes.searchsorted(low, "left")
+                    end = indexes.searchsorted(up, "right")
+                    start = max(start - 1, 0)
+                    end = min(end + 1, len(indexes))
+                    indexes_between = indexes[start:end].to_list()
+                    indexes_between_ranges.append(indexes_between)
+                else:
+                    start = indexes.searchsorted(low, "left")
+                    end = indexes.searchsorted(up, "right")
+                    indexes_between = indexes[start:end].to_list()
+                    indexes_between_ranges.append(indexes_between)
             else:
-                indexes_between = [i for i in indexes if low <= i <= up]
-                indexes_between_ranges.append(indexes_between)
+                if method == "surrounding":
+                    start = indexes.index(low)
+                    end = indexes.index(up)
+                    start = max(start - 1, 0)
+                    end = min(end + 1, len(indexes))
+                    indexes_between = indexes[start:end]
+                    indexes_between_ranges.append(indexes_between)
+                else:
+                    indexes_between = [i for i in indexes if low <= i <= up]
+                    indexes_between_ranges.append(indexes_between)
         return indexes_between_ranges
 
     @staticmethod

polytope/datacube/index_tree.py

@@ -58,7 +58,18 @@ def __hash__(self):
     def __eq__(self, other):
         if not isinstance(other, IndexTree):
             return False
-        return (self.axis.name, self.value) == (other.axis.name, other.value)
+        if self.axis.name != other.axis.name:
+            return False
+        else:
+            if other.value == self.value:
+                return True
+            if other.value - 2 * other.axis.tol <= self.value <= other.value + 2 * other.axis.tol:
+                return True
+            elif self.value - 2 * self.axis.tol <= other.value <= self.value + 2 * self.axis.tol:
+                return True
+            else:
+                return False
+        # return (self.axis.name, self.value) == (other.axis.name, other.value)
 
     def __lt__(self, other):
         return (self.axis.name, self.value) < (other.axis.name, other.value)
@@ -73,12 +84,12 @@ def add_child(self, node):
         self.children.add(node)
         node._parent = self
 
-    def create_child(self, axis, value):
+    def create_child_not_safe(self, axis, value):
         node = IndexTree(axis, value)
         self.add_child(node)
         return node
 
-    def create_child_safe(self, axis, value):
+    def create_child(self, axis, value):
         node = IndexTree(axis, value)
         existing = self.find_child(node)
         if not existing:

polytope/engine/hullslicer.py

@@ -1,3 +1,4 @@
+import math
 from copy import copy
 from itertools import chain
 from typing import List
@@ -46,12 +47,19 @@ def _build_sliceable_child(self, polytope, ax, node, datacube, lower, upper, nex
         lower = ax.from_float(lower - tol)
         upper = ax.from_float(upper + tol)
         flattened = node.flatten()
-        for value in datacube.get_indices(flattened, ax, lower, upper):
+        method = polytope.method
+        for value in datacube.get_indices(flattened, ax, lower, upper, method):
             # convert to float for slicing
             fvalue = ax.to_float(value)
             new_polytope = slice(polytope, ax.name, fvalue)
             # store the native type
-            child = node.create_child(ax, value)
+            # remapped_val = (ax.remap([value, value])[0][0] + ax.remap([value, value])[0][1])/2
+            remapped_val = value
+            if ax.is_cyclic:
+                remapped_val = (ax.remap([value, value])[0][0] + ax.remap([value, value])[0][1]) / 2
+                remapped_val = round(remapped_val, int(-math.log10(ax.tol)))
+            # child = node.create_child(ax, value)
+            child = node.create_child(ax, remapped_val)
             child["unsliced_polytopes"] = copy(node["unsliced_polytopes"])
             child["unsliced_polytopes"].remove(polytope)
             if new_polytope is not None:

polytope/polytope.py

@@ -42,7 +42,7 @@ def slice(self, polytopes: List[ConvexPolytope]):
         """Low-level API which takes a polytope geometry object and uses it to slice the datacube"""
         return self.engine.extract(self.datacube, polytopes)
 
-    def retrieve(self, request: Request):
+    def retrieve(self, request: Request, method="standard"):
         """Higher-level API which takes a request and uses it to slice the datacube"""
         request_tree = self.engine.extract(self.datacube, request.polytopes())
         self.datacube.get(request_tree)

polytope/shapes.py

@@ -23,9 +23,10 @@ def axes(self) -> List[str]:
 
 
 class ConvexPolytope(Shape):
-    def __init__(self, axes, points):
+    def __init__(self, axes, points, method=None):
         self._axes = list(axes)
         self.points = points
+        self.method = method
 
     def extents(self, axis):
         slice_axis_idx = self.axes().index(axis)
@@ -48,15 +49,16 @@ def polytope(self):
 class Select(Shape):
     """Matches several discrete value"""
 
-    def __init__(self, axis, values):
+    def __init__(self, axis, values, method=None):
         self.axis = axis
         self.values = values
+        self.method = method
 
     def axes(self):
         return [self.axis]
 
     def polytope(self):
-        return [ConvexPolytope([self.axis], [[v]]) for v in self.values]
+        return [ConvexPolytope([self.axis], [[v]], self.method) for v in self.values]
 
 
 class Span(Shape):