diff --git a/polytope_feature/datacube/tensor_index_tree.py b/polytope_feature/datacube/tensor_index_tree.py index 604aa744..efc35c3e 100644 --- a/polytope_feature/datacube/tensor_index_tree.py +++ b/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 diff --git a/polytope_feature/engine/hullslicer.py b/polytope_feature/engine/hullslicer.py index 8d3df84a..d596004c 100644 --- a/polytope_feature/engine/hullslicer.py +++ b/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) diff --git a/polytope_feature/shapes.py b/polytope_feature/shapes.py index c3a45a6a..a8a0f630 100644 --- a/polytope_feature/shapes.py +++ b/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) diff --git a/tests/test_shape_labels.py b/tests/test_shape_labels.py new file mode 100644 index 00000000..e5e4619d --- /dev/null +++ b/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"]