diff --git a/.gitignore b/.gitignore index 037872e66..7b8e19f62 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ polytope.egg-info .pytest_cache *.prof -*.idx \ No newline at end of file +*.idx +*.grib \ No newline at end of file diff --git a/polytope/datacube/datacube_axis.py b/polytope/datacube/datacube_axis.py index 77a223710..625051f7c 100644 --- a/polytope/datacube/datacube_axis.py +++ b/polytope/datacube/datacube_axis.py @@ -1,3 +1,4 @@ +import math from abc import ABC, abstractmethod from copy import deepcopy from typing import Any, List @@ -18,6 +19,10 @@ def update_range(): def to_intervals(range): update_range() + if range[0] == -math.inf: + range[0] = cls.range[0] + if range[1] == math.inf: + range[1] = cls.range[1] axis_lower = cls.range[0] axis_upper = cls.range[1] axis_range = axis_upper - axis_lower diff --git a/polytope/polytope.py b/polytope/polytope.py index 5e8fc7133..f6d4a723e 100644 --- a/polytope/polytope.py +++ b/polytope/polytope.py @@ -29,6 +29,12 @@ def polytopes(self): polytopes.extend(shape.polytope()) return polytopes + def __repr__(self): + return_str = "" + for shape in self.shapes: + return_str += shape.__repr__() + "\n" + return return_str + class Polytope: def __init__(self, datacube, engine=None, axis_options={}): diff --git a/polytope/shapes.py b/polytope/shapes.py index 94562abc5..cc77273a7 100644 --- a/polytope/shapes.py +++ b/polytope/shapes.py @@ -60,11 +60,36 @@ def axes(self): def polytope(self): return [ConvexPolytope([self.axis], [[v]], self.method) for v in self.values] + def __repr__(self): + return f"Select in {self.axis} with points {self.values}" + + +class Point(Shape): + """Matches several discrete value""" + + def __init__(self, axes, values, method=None): + self._axes = axes + self.values = values + self.method = method + self.polytopes = [] + for i in range(len(axes)): + polytope_points = [v[i] for v in self.values] + self.polytopes.append(ConvexPolytope([axes[i]], [polytope_points], method)) + + def axes(self): + return self._axes + + def polytope(self): + return self.polytopes + + def __repr__(self): + return f"Point in {self._axes} with points {self.values}" + class Span(Shape): """1-D range along a single axis""" - def __init__(self, axis, lower=None, upper=None): + def __init__(self, axis, lower=-math.inf, upper=math.inf): assert not isinstance(lower, list) assert not isinstance(upper, list) self.axis = axis @@ -77,6 +102,9 @@ def axes(self): def polytope(self): return [ConvexPolytope([self.axis], [[self.lower], [self.upper]])] + def __repr__(self): + return f"Span in {self.axis} with range from {self.lower} to {self.upper}" + class All(Span): """Matches all indices in an axis""" @@ -84,12 +112,17 @@ class All(Span): def __init__(self, axis): super().__init__(axis) + def __repr__(self): + return f"All in {self.axis}" + 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): dimension = len(axes) + self._lower_corner = lower_corner + self._upper_corner = upper_corner self._axes = axes assert len(lower_corner) == dimension assert len(upper_corner) == dimension @@ -121,6 +154,9 @@ def axes(self): def polytope(self): return [ConvexPolytope(self.axes(), self.vertices)] + def __repr__(self): + return f"Box in {self._axes} with with lower corner {self._lower_corner} and upper corner{self._upper_corner}" + class Disk(Shape): """2-D shape bounded by an ellipse""" @@ -159,6 +195,9 @@ def axes(self): def polytope(self): return [ConvexPolytope(self.axes(), self.points)] + def __repr__(self): + return f"Disk in {self._axes} with centred at {self.centre} and with radius {self.radius}" + class Ellipsoid(Shape): # Here we use the formula for the inscribed circle in an icosahedron @@ -211,12 +250,18 @@ def _icosahedron_edge_length_coeff(self): def polytope(self): return [ConvexPolytope(self.axes(), self.points)] + def __repr__(self): + return f"Ellipsoid in {self._axes} with centred at {self.centre} and with radius {self.radius}" + 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): self._axes = axes + self._start = start + self._end = end + self._shape = shape assert shape.axes() == self.axes() assert len(start) == len(self.axes()) @@ -240,12 +285,18 @@ def axes(self): def polytope(self): return self.polytopes + def __repr__(self): + return f"PathSegment in {self._axes} obtained by sweeping a {self._shape.__repr__()} \ + between the points {self._start} and {self._end}" + 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): self._axes = axes + self._shape = shape + self._points = points assert shape.axes() == self.axes() for p in points: @@ -267,6 +318,10 @@ def axes(self): def polytope(self): return self.union.polytope() + def __repr__(self): + return f"Path in {self._axes} obtained by sweeping a {self._shape.__repr__()} \ + between the points {self._points}" + class Union(Shape): """N-D union of two shapes with the same axes""" @@ -277,6 +332,7 @@ def __init__(self, axes, *shapes): assert s.axes() == self.axes() self.polytopes = [] + self._shapes = shapes for s in shapes: self.polytopes.extend(s.polytope()) @@ -287,6 +343,9 @@ def axes(self): def polytope(self): return self.polytopes + def __repr__(self): + return f"Union in {self._axes} of the shapes {self._shapes}" + class Polygon(Shape): """2-D polygon defined by a set of exterior points""" @@ -297,6 +356,7 @@ def __init__(self, axes, points): for p in points: assert len(p) == 2 + self._points = points triangles = tripy.earclip(points) self.polytopes = [] @@ -313,3 +373,6 @@ def axes(self): def polytope(self): return self.polytopes + + def __repr__(self): + return f"Polygon in {self._axes} with points {self._points}" diff --git a/tests/test_point_shape.py b/tests/test_point_shape.py new file mode 100644 index 000000000..0bc203d61 --- /dev/null +++ b/tests/test_point_shape.py @@ -0,0 +1,36 @@ +import numpy as np +import pandas as pd +import xarray as xr + +from polytope.datacube.backends.xarray import XArrayDatacube +from polytope.engine.hullslicer import HullSlicer +from polytope.polytope import Polytope, Request +from polytope.shapes import Point, 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() + self.API = Polytope(datacube=array, engine=self.slicer) + + def test_point(self): + request = Request(Point(["step", "level"], [[3, 10]]), Select("date", ["2000-01-01"])) + result = self.API.retrieve(request) + assert len(result.leaves) == 1 + assert result.leaves[0].axis.name == "level" + + def test_point_surrounding_step(self): + request = Request(Point(["step", "level"], [[2, 10]], method="surrounding"), Select("date", ["2000-01-01"])) + result = self.API.retrieve(request) + assert len(result.leaves) == 6 diff --git a/tests/test_shapes.py b/tests/test_shapes.py new file mode 100644 index 000000000..ba355dd4c --- /dev/null +++ b/tests/test_shapes.py @@ -0,0 +1,74 @@ +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from polytope.datacube.backends.FDB_datacube import FDBDatacube +from polytope.datacube.backends.xarray import XArrayDatacube +from polytope.engine.hullslicer import HullSlicer +from polytope.polytope import Polytope, Request +from polytope.shapes import All, Select, Span + + +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, 360), + dims=("date", "step", "level", "longitude"), + coords={ + "date": pd.date_range("2000-01-01", "2000-01-03", 3), + "step": [0, 3, 6, 9, 12, 15], + "level": range(1, 130), + "longitude": range(0, 360), + }, + ) + self.xarraydatacube = XArrayDatacube(array) + self.options = {"longitude": {"transformation": {"cyclic": [0, 360]}}} + self.slicer = HullSlicer() + self.API = Polytope(datacube=array, engine=self.slicer, axis_options=self.options) + + def test_all(self): + request = Request(Select("step", [3]), Select("date", ["2000-01-01"]), All("level"), Select("longitude", [1])) + result = self.API.retrieve(request) + assert len(result.leaves) == 129 + + def test_all_cyclic(self): + request = Request(Select("step", [3]), Select("date", ["2000-01-01"]), Select("level", [1]), All("longitude")) + result = self.API.retrieve(request) + # result.pprint() + assert len(result.leaves) == 360 + + @pytest.mark.skip(reason="can't install fdb branch on CI") + def test_all_mapper_cyclic(self): + self.options = { + "values": { + "transformation": { + "mapper": {"type": "octahedral", "resolution": 1280, "axes": ["latitude", "longitude"]} + } + }, + "date": {"transformation": {"merge": {"with": "time", "linkers": ["T", "00"]}}}, + "step": {"transformation": {"type_change": "int"}}, + "longitude": {"transformation": {"cyclic": [0, 360]}}, + } + self.config = {"class": "od", "expver": "0001", "levtype": "sfc", "step": 11} + self.fdbdatacube = FDBDatacube(self.config, axis_options=self.options) + self.slicer = HullSlicer() + self.API = Polytope(datacube=self.fdbdatacube, engine=self.slicer, axis_options=self.options) + + request = Request( + Select("step", [11]), + Select("levtype", ["sfc"]), + Select("date", [pd.Timestamp("20230710T120000")]), + Select("domain", ["g"]), + Select("expver", ["0001"]), + Select("param", ["151130"]), + Select("class", ["od"]), + Select("stream", ["oper"]), + Select("type", ["fc"]), + Span("latitude", 89.9, 90), + All("longitude"), + ) + result = self.API.retrieve(request) + # result.pprint() + assert len(result.leaves) == 20