add labels to shapes

polytope_feature/datacube/tensor_index_tree.py

@@ -34,6 +34,7 @@ def __init__(self, axis=root, values=tuple()):
         self.ancestors = []
         self.indexes = []
         self.hidden = False
+        self.labels = []
 
     @property
     def leaves(self):
@@ -108,12 +109,14 @@ def add_value(self, value):
         new_values.sort()
         self.values = tuple(new_values)
 
-    def create_child(self, axis, value, next_nodes):
+    def create_child(self, axis, value, next_nodes, polytope_label=None):
         node = TensorIndexTree(axis, (value,))
         existing_child = self.find_child(node)
         if not existing_child:
             self.add_child(node)
+            node.labels.append(polytope_label)
             return (node, next_nodes)
+        existing_child.labels.append(polytope_label)
         return (existing_child, next_nodes)
 
     @property

polytope_feature/engine/hullslicer.py

@@ -55,7 +55,7 @@ def _build_unsliceable_child(self, polytope, ax, node, datacube, lowers, next_no
 
             if datacube_has_index:
                 if i == 0:
-                    (child, next_nodes) = node.create_child(ax, lower, next_nodes)
+                    (child, next_nodes) = node.create_child(ax, lower, next_nodes, polytope.label)
                     child["unsliced_polytopes"] = copy(node["unsliced_polytopes"])
                     child["unsliced_polytopes"].remove(polytope)
                     next_nodes.append(child)
@@ -114,7 +114,8 @@ def _build_sliceable_child(self, polytope, ax, node, datacube, values, next_node
                 fvalue = ax.to_float(value)
                 new_polytope = slice(polytope, ax.name, fvalue, slice_axis_idx)
                 remapped_val = self.remap_values(ax, value)
-                (child, next_nodes) = node.create_child(ax, remapped_val, next_nodes)
+                polytope_label = polytope.label
+                (child, next_nodes) = node.create_child(ax, remapped_val, next_nodes, polytope_label)
                 child["unsliced_polytopes"] = copy(node["unsliced_polytopes"])
                 child["unsliced_polytopes"].remove(polytope)
                 if new_polytope is not None:
@@ -295,7 +296,7 @@ def slice(polytope: ConvexPolytope, axis, value, slice_axis_idx):
     axes.remove(axis)
 
     if len(intersects) < len(intersects[0]) + 1:
-        return ConvexPolytope(axes, intersects)
+        return ConvexPolytope(axes, intersects, label=polytope.label)
     # Compute convex hull (removing interior points)
     if len(intersects[0]) == 0:
         return None
@@ -310,6 +311,6 @@ def slice(polytope: ConvexPolytope, axis, value, slice_axis_idx):
 
         except scipy.spatial.qhull.QhullError as e:
             if "less than" or "flat" in str(e):
-                return ConvexPolytope(axes, intersects)
+                return ConvexPolytope(axes, intersects, label=polytope.label)
     # Sliced result is simply the convex hull
-    return ConvexPolytope(axes, [intersects[i] for i in vertices])
+    return ConvexPolytope(axes, [intersects[i] for i in vertices], label=polytope.label)

polytope_feature/shapes.py

@@ -23,7 +23,7 @@ def axes(self) -> List[str]:
 
 
 class ConvexPolytope(Shape):
-    def __init__(self, axes, points, method=None, is_orthogonal=False):
+    def __init__(self, axes, points, method=None, is_orthogonal=False, label=None):
         self._axes = list(axes)
         self.is_flat = False
         if len(self._axes) == 1:
@@ -32,6 +32,7 @@ def __init__(self, axes, points, method=None, is_orthogonal=False):
         self.method = method
         self.is_orthogonal = is_orthogonal
         self.is_in_union = False
+        self.label = label
 
     def add_to_union(self):
         self.is_in_union = True
@@ -62,16 +63,19 @@ def polytope(self):
 class Select(Shape):
     """Matches several discrete value"""
 
-    def __init__(self, axis, values, method=None):
+    def __init__(self, axis, values, method=None, label=None):
         self.axis = axis
         self.values = values
         self.method = method
+        self.label = label
 
     def axes(self):
         return [self.axis]
 
     def polytope(self):
-        return [ConvexPolytope([self.axis], [[v]], self.method, is_orthogonal=True) for v in self.values]
+        return [
+            ConvexPolytope([self.axis], [[v]], self.method, is_orthogonal=True, label=self.label) for v in self.values
+        ]
 
     def __repr__(self):
         return f"Select in {self.axis} with points {self.values}"
@@ -80,17 +84,21 @@ def __repr__(self):
 class Point(Shape):
     """Matches several discrete value"""
 
-    def __init__(self, axes, values, method=None):
+    def __init__(self, axes, values, method=None, label=None):
         self._axes = axes
         self.values = values
         self.method = method
         self.polytopes = []
+        self.label = label
         if method == "nearest":
             assert len(self.values) == 1
         for i in range(len(axes)):
             polytope_points = [v[i] for v in self.values]
             self.polytopes.extend(
-                [ConvexPolytope([axes[i]], [[point]], self.method, is_orthogonal=True) for point in polytope_points]
+                [
+                    ConvexPolytope([axes[i]], [[point]], self.method, is_orthogonal=True, label=self.label)
+                    for point in polytope_points
+                ]
             )
 
     def axes(self):
@@ -106,18 +114,19 @@ def __repr__(self):
 class Span(Shape):
     """1-D range along a single axis"""
 
-    def __init__(self, axis, lower=-math.inf, upper=math.inf):
+    def __init__(self, axis, lower=-math.inf, upper=math.inf, label=None):
         assert not isinstance(lower, list)
         assert not isinstance(upper, list)
         self.axis = axis
         self.lower = lower
         self.upper = upper
+        self.label = label
 
     def axes(self):
         return [self.axis]
 
     def polytope(self):
-        return [ConvexPolytope([self.axis], [[self.lower], [self.upper]], is_orthogonal=True)]
+        return [ConvexPolytope([self.axis], [[self.lower], [self.upper]], is_orthogonal=True, label=self.label)]
 
     def __repr__(self):
         return f"Span in {self.axis} with range from {self.lower} to {self.upper}"
@@ -136,11 +145,12 @@ def __repr__(self):
 class Box(Shape):
     """N-D axis-aligned bounding box (AABB), specified by two opposite corners"""
 
-    def __init__(self, axes, lower_corner=None, upper_corner=None):
+    def __init__(self, axes, lower_corner=None, upper_corner=None, label=None):
         dimension = len(axes)
         self._lower_corner = lower_corner
         self._upper_corner = upper_corner
         self._axes = axes
+        self.label = label
         assert len(lower_corner) == dimension
         assert len(upper_corner) == dimension
 
@@ -168,7 +178,7 @@ def axes(self):
         return self._axes
 
     def polytope(self):
-        return [ConvexPolytope(self.axes(), self.vertices, is_orthogonal=True)]
+        return [ConvexPolytope(self.axes(), self.vertices, is_orthogonal=True, label=self.label)]
 
     def __repr__(self):
         return f"Box in {self._axes} with with lower corner {self._lower_corner} and upper corner{self._upper_corner}"
@@ -180,11 +190,12 @@ class Disk(Shape):
     # NB radius is two dimensional
     # NB number of segments is hard-coded, not exposed to user
 
-    def __init__(self, axes, centre=[0, 0], radius=[1, 1]):
+    def __init__(self, axes, centre=[0, 0], radius=[1, 1], label=None):
         self._axes = axes
         self.centre = centre
         self.radius = radius
         self.segments = 12
+        self.label = label
 
         assert len(axes) == 2
         assert len(centre) == 2
@@ -209,7 +220,7 @@ def axes(self):
         return self._axes
 
     def polytope(self):
-        return [ConvexPolytope(self.axes(), self.points)]
+        return [ConvexPolytope(self.axes(), self.points, label=self.label)]
 
     def __repr__(self):
         return f"Disk in {self._axes} with centred at {self.centre} and with radius {self.radius}"
@@ -219,10 +230,11 @@ class Ellipsoid(Shape):
     # Here we use the formula for the inscribed circle in an icosahedron
     # See https://en.wikipedia.org/wiki/Platonic_solid
 
-    def __init__(self, axes, centre=[0, 0, 0], radius=[1, 1, 1]):
+    def __init__(self, axes, centre=[0, 0, 0], radius=[1, 1, 1], label=None):
         self._axes = axes
         self.centre = centre
         self.radius = radius
+        self.label = label
 
         assert len(axes) == 3
         assert len(centre) == 3
@@ -264,7 +276,7 @@ def _icosahedron_edge_length_coeff(self):
         return edge_length
 
     def polytope(self):
-        return [ConvexPolytope(self.axes(), self.points)]
+        return [ConvexPolytope(self.axes(), self.points, label=self.label)]
 
     def __repr__(self):
         return f"Ellipsoid in {self._axes} with centred at {self.centre} and with radius {self.radius}"
@@ -273,11 +285,12 @@ def __repr__(self):
 class PathSegment(Shape):
     """N-D polytope defined by a shape which is swept along a straight line between two points"""
 
-    def __init__(self, axes, shape: Shape, start: List, end: List):
+    def __init__(self, axes, shape: Shape, start: List, end: List, label=None):
         self._axes = axes
         self._start = start
         self._end = end
         self._shape = shape
+        self.label = label
 
         assert shape.axes() == self.axes()
         assert len(start) == len(self.axes())
@@ -293,7 +306,7 @@ def __init__(self, axes, shape: Shape, start: List, end: List):
             for p in polytope.points:
                 points.append([a + b for a, b in zip(p, start)])
                 points.append([a + b for a, b in zip(p, end)])
-            poly = ConvexPolytope(self.axes(), points)
+            poly = ConvexPolytope(self.axes(), points, label=self.label)
             poly.add_to_union()
             self.polytopes.append(poly)
 
@@ -311,10 +324,11 @@ def __repr__(self):
 class Path(Shape):
     """N-D polytope defined by a shape which is swept along a polyline defined by multiple points"""
 
-    def __init__(self, axes, shape, *points, closed=False):
+    def __init__(self, axes, shape, *points, closed=False, label=None):
         self._axes = axes
         self._shape = shape
         self._points = points
+        self.label = label
 
         assert shape.axes() == self.axes()
         for p in points:
@@ -323,10 +337,10 @@ def __init__(self, axes, shape, *points, closed=False):
         path_segments = []
 
         for i in range(0, len(points) - 1):
-            path_segments.append(PathSegment(axes, shape, points[i], points[i + 1]))
+            path_segments.append(PathSegment(axes, shape, points[i], points[i + 1], label=self.label))
 
         if closed:
-            path_segments.append(PathSegment(axes, shape, points[-1], points[0]))
+            path_segments.append(PathSegment(axes, shape, points[-1], points[0], label=self.label))
 
         self.union = Union(self.axes(), *path_segments)
 
@@ -370,23 +384,25 @@ def __repr__(self):
 class Polygon(Shape):
     """2-D polygon defined by a set of exterior points"""
 
-    def __init__(self, axes, points):
+    def __init__(self, axes, points, label=None):
         self._axes = axes
         assert len(axes) == 2
         for p in points:
             assert len(p) == 2
 
+        self.label = label
+
         self._points = points
         triangles = tripy.earclip(points)
         self.polytopes = []
 
         if len(points) > 0 and len(triangles) == 0:
-            self.polytopes = [ConvexPolytope(self.axes(), points)]
+            self.polytopes = [ConvexPolytope(self.axes(), points, label=self.label)]
 
         else:
             for t in triangles:
                 tri_points = [list(point) for point in t]
-                poly = ConvexPolytope(self.axes(), tri_points)
+                poly = ConvexPolytope(self.axes(), tri_points, label=self.label)
                 poly.add_to_union()
                 self.polytopes.append(poly)
 

tests/test_shape_labels.py

@@ -0,0 +1,38 @@
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+from polytope_feature.datacube.backends.xarray import XArrayDatacube
+from polytope_feature.engine.hullslicer import HullSlicer
+from polytope_feature.polytope import Polytope, Request
+from polytope_feature.shapes import Box, Select
+
+
+class TestSlicing3DXarrayDatacube:
+    def setup_method(self, method):
+        # Create a dataarray with 3 labelled axes using different index types
+        array = xr.DataArray(
+            np.random.randn(3, 6, 129),
+            dims=("date", "step", "level"),
+            coords={
+                "date": pd.date_range("2000-01-01", "2000-01-03", 3),
+                "step": [0, 3, 6, 9, 12, 15],
+                "level": range(1, 130),
+            },
+        )
+        self.xarraydatacube = XArrayDatacube(array)
+        self.slicer = HullSlicer()
+        options = {"compressed_axes_config": ["date", "step", "level"]}
+        self.API = Polytope(datacube=array, engine=self.slicer, options=options)
+
+    # Testing different shapes
+
+    def test_2D_box(self):
+        request = Request(
+            Box(["step", "level"], [3, 10], [6, 11], label="box1"), Select("date", ["2000-01-01"], label="select1")
+        )
+        result = self.API.retrieve(request)
+        assert len(result.leaves) == 1
+        assert result.leaves[0].labels == ["box1"]
+        assert result.leaves[0].parent.labels == ["box1"]
+        assert result.leaves[0].parent.parent.labels == ["select1"]