Merge pull request #352 from ecmwf/develop

v1.0.36

Author: mathleur

.gitignore

@@ -27,4 +27,5 @@ newest-polytope-venv
 serializedTree
 new_polytope_venv
 *.json
-venv_python3_11
+venv_python3_11
+tests/data

polytope_feature/datacube/backends/fdb.py

@@ -227,17 +227,22 @@ def nearest_lat_lon_search(self, requests):
             first_ax_name = requests.children[0].axis.name
             second_ax_name = requests.children[0].children[0].axis.name
 
-            if first_ax_name not in self.nearest_search.keys() or second_ax_name not in self.nearest_search.keys():
+            axes_in_nearest_search = [
+                first_ax_name not in self.nearest_search.keys(),
+                second_ax_name not in self.nearest_search.keys(),
+            ]
+
+            if all(not item for item in axes_in_nearest_search):
                 raise Exception("nearest point search axes are wrong")
 
             second_ax = requests.children[0].children[0].axis
+            nearest_pts = self.nearest_search.get(first_ax_name, None)
+            if nearest_pts is None:
+                nearest_pts = self.nearest_search[second_ax_name]
 
-            nearest_pts = [
-                [lat_val, second_ax._remap_val_to_axis_range(lon_val)]
-                for (lat_val, lon_val) in zip(
-                    self.nearest_search[first_ax_name][0], self.nearest_search[second_ax_name][0]
-                )
-            ]
+            transformed_nearest_pts = []
+            for point in nearest_pts:
+                transformed_nearest_pts.append([point[0], second_ax._remap_val_to_axis_range(point[1])])
 
             found_latlon_pts = []
             for lat_child in requests.children:
@@ -246,7 +251,7 @@ def nearest_lat_lon_search(self, requests):
 
             # now find the nearest lat lon to the points requested
             nearest_latlons = []
-            for pt in nearest_pts:
+            for pt in transformed_nearest_pts:
                 nearest_latlon = nearest_pt(found_latlon_pts, pt)
                 nearest_latlons.append(nearest_latlon)
 

polytope_feature/engine/hullslicer.py

@@ -8,7 +8,7 @@
 from ..datacube.backends.datacube import Datacube
 from ..datacube.datacube_axis import UnsliceableDatacubeAxis
 from ..datacube.tensor_index_tree import TensorIndexTree
-from ..shapes import ConvexPolytope
+from ..shapes import ConvexPolytope, Product
 from ..utility.combinatorics import group, tensor_product
 from ..utility.exceptions import UnsliceableShapeError
 from ..utility.geometry import lerp
@@ -76,8 +76,6 @@ def find_values_between(self, polytope, ax, node, datacube, lower, upper):
         upper = ax.from_float(upper + tol)
         flattened = node.flatten()
         method = polytope.method
-        if method == "nearest":
-            datacube.nearest_search[ax.name] = polytope.points
 
         # NOTE: caching
         # Create a coupled_axes list inside of datacube and add to it during axis formation, then here
@@ -214,7 +212,11 @@ def extract(self, datacube: Datacube, polytopes: List[ConvexPolytope]):
 
         # Convert the polytope points to float type to support triangulation and interpolation
         for p in polytopes:
-            self._unique_continuous_points(p, datacube)
+            if isinstance(p, Product):
+                for poly in p.polytope():
+                    self._unique_continuous_points(poly, datacube)
+            else:
+                self._unique_continuous_points(p, datacube)
 
         groups, input_axes = group(polytopes)
         datacube.validate(input_axes)
@@ -233,7 +235,16 @@ def extract(self, datacube: Datacube, polytopes: List[ConvexPolytope]):
                     new_c.extend(combi)
                 else:
                     new_c.append(combi)
-            r["unsliced_polytopes"] = set(new_c)
+            # NOTE TODO: here some of the polys in new_c can be a Product shape instead of a ConvexPolytope
+            # -> need to go through the polytopes in new_c and replace the Products with their sub-ConvexPolytopes
+            final_polys = []
+            for poly in new_c:
+                if isinstance(poly, Product):
+                    final_polys.extend(poly.polytope())
+                else:
+                    final_polys.append(poly)
+            # r["unsliced_polytopes"] = set(new_c)
+            r["unsliced_polytopes"] = set(final_polys)
             current_nodes = [r]
             for ax in datacube.axes.values():
                 next_nodes = []

polytope_feature/polytope.py

@@ -64,6 +64,13 @@ def retrieve(self, request: Request, method="standard"):
         """Higher-level API which takes a request and uses it to slice the datacube"""
         logging.info("Starting request for %s ", self.context)
         self.datacube.check_branching_axes(request)
+        for polytope in request.polytopes():
+            method = polytope.method
+            if method == "nearest":
+                if self.datacube.nearest_search.get(polytope.axes()[0], None) is None:
+                    self.datacube.nearest_search[polytope.axes()[0]] = polytope.values
+                else:
+                    self.datacube.nearest_search[polytope.axes()[0]].append(polytope.values[0])
         request_tree = self.engine.extract(self.datacube, request.polytopes())
         logging.info("Created request tree for %s ", self.context)
         self.datacube.get(request_tree, self.context)

polytope_feature/shapes.py

@@ -61,6 +61,42 @@ def polytope(self):
         return [self]
 
 
+class Product(Shape):
+    """Shape that takes two polytopes and 'multiplies' them together to obtain higher-dimensional shape"""
+
+    def __init__(self, *polytopes, method, value):
+        # TODO
+        all_axes = []
+        for poly in polytopes:
+            all_axes.extend(poly.axes())
+        self._axes = list(set(all_axes))
+        # Check there weren't any duplicates in the polytopes' axes
+        assert len(self._axes) == len(all_axes)
+
+        self._polytopes = []
+        for poly in polytopes:
+            self._polytopes.extend(poly.polytope())
+
+        self.is_in_union = False
+        self.method = method
+        self.values = value
+
+        self.is_orthogonal = False
+
+        polys_orthogonal = [poly.is_orthogonal for poly in polytopes]
+        if all(polys_orthogonal):
+            self.is_orthogonal = True
+
+    def add_to_union(self):
+        self.is_in_union = True
+
+    def axes(self):
+        return self._axes
+
+    def polytope(self):
+        return self._polytopes
+
+
 # This is the only shape which can slice on axes without a discretizer or interpolator
 class Select(Shape):
     """Matches several discrete values"""
@@ -89,19 +125,22 @@ def __init__(self, axes, values, method=None):
         self._axes = axes
         self.values = values
         self.method = method
-        self.polytopes = []
-        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]
-            )
+        assert len(values) == 1
 
     def axes(self):
         return self._axes
 
     def polytope(self):
+        # TODO: change this to use the Product instead and return a Product here of the two 1D selects
+
+        polytopes = []
+        for point in self.values:
+            poly_to_mult = []
+            for i in range(len(self._axes)):
+                poly_to_mult.append(ConvexPolytope([self._axes[i]], [[point[i]]], self.method, is_orthogonal=True))
+            polytopes.append(Product(*poly_to_mult, method=self.method, value=[point]))
+        self.polytopes = polytopes
+
         return self.polytopes
 
     def __repr__(self):

polytope_feature/version.py

@@ -1 +1 @@
-__version__ = "1.0.35"
+__version__ = "1.0.36"

tests/test_point_shape.py

@@ -4,7 +4,7 @@
 
 from polytope_feature.engine.hullslicer import HullSlicer
 from polytope_feature.polytope import Polytope, Request
-from polytope_feature.shapes import Point, Select
+from polytope_feature.shapes import Point, Select, Union
 
 
 class TestSlicing3DXarrayDatacube:
@@ -30,20 +30,26 @@ def test_point(self):
         assert result.leaves[0].axis.name == "level"
 
     def test_multiple_points(self):
-        request = Request(Point(["step", "level"], [[3, 10], [3, 12]]), Select("date", ["2000-01-01"]))
+        # request = Request(Point(["step", "level"], [[3, 10], [3, 12]]), Select("date", ["2000-01-01"]))
+        request = Request(
+            Union(["step", "level"], Point(["step", "level"], [[3, 10]]), Point(["step", "level"], [[3, 12]])),
+            Select("date", ["2000-01-01"]),
+        )
         result = self.API.retrieve(request)
         result.pprint()
-        assert len(result.leaves) == 1
+        assert len(result.leaves) == 2
         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)
+        result.pprint()
         assert len(result.leaves) == 1
         assert np.shape(result.leaves[0].result[1]) == (1, 2, 3)
 
     def test_point_surrounding_exact_step(self):
         request = Request(Point(["step", "level"], [[3, 10]], method="surrounding"), Select("date", ["2000-01-01"]))
         result = self.API.retrieve(request)
+        result.pprint()
         assert len(result.leaves) == 1
         assert np.shape(result.leaves[0].result[1]) == (1, 3, 3)

tests/test_point_union.py

@@ -0,0 +1,148 @@
+import pandas as pd
+import pytest
+
+from polytope_feature.engine.hullslicer import HullSlicer
+from polytope_feature.polytope import Polytope, Request
+from polytope_feature.shapes import Point, Select, Span, Union
+
+
+class TestSlicingFDBDatacube:
+    def setup_method(self, method):
+        # Create a dataarray with 3 labelled axes using different index types
+        self.options = {
+            "axis_config": [
+                {"axis_name": "number", "transformations": [{"name": "type_change", "type": "int"}]},
+                {"axis_name": "step", "transformations": [{"name": "type_change", "type": "int"}]},
+                {
+                    "axis_name": "date",
+                    "transformations": [{"name": "merge", "other_axis": "time", "linkers": ["T", "00"]}],
+                },
+                {
+                    "axis_name": "values",
+                    "transformations": [
+                        {"name": "mapper", "type": "octahedral", "resolution": 1280, "axes": ["latitude", "longitude"]}
+                    ],
+                },
+                {"axis_name": "latitude", "transformations": [{"name": "reverse", "is_reverse": True}]},
+                {"axis_name": "longitude", "transformations": [{"name": "cyclic", "range": [0, 360]}]},
+            ],
+            "pre_path": {"class": "od", "expver": "0001", "levtype": "sfc", "stream": "oper"},
+            "compressed_axes_config": [
+                "longitude",
+                "latitude",
+                "levtype",
+                "step",
+                "date",
+                "domain",
+                "expver",
+                "param",
+                "class",
+                "stream",
+                "type",
+            ],
+        }
+
+    # Testing different shapes
+    @pytest.mark.fdb
+    def test_fdb_datacube(self):
+        import pygribjump as gj
+
+        request = Request(
+            Select("step", [0]),
+            Select("levtype", ["sfc"]),
+            Span("date", pd.Timestamp("20230625T120000"), pd.Timestamp("20230626T120000")),
+            Select("domain", ["g"]),
+            Select("expver", ["0001"]),
+            Select("param", ["167"]),
+            Select("class", ["od"]),
+            Select("stream", ["oper"]),
+            Select("type", ["an"]),
+            Union(
+                ["latitude", "longitude"],
+                Point(["latitude", "longitude"], [[20, 20]], method="nearest"),
+                Point(["latitude", "longitude"], [[0, 0]], method="nearest"),
+                Point(["latitude", "longitude"], [[0, 20]], method="nearest"),
+                Point(["latitude", "longitude"], [[25, 30]], method="nearest"),
+                Point(["latitude", "longitude"], [[-30, 90]], method="nearest"),
+                Point(["latitude", "longitude"], [[-60, -30]], method="nearest"),
+                Point(["latitude", "longitude"], [[-15, -45]], method="nearest"),
+                Point(["latitude", "longitude"], [[20, 0]], method="nearest"),
+            ),
+        )
+
+        self.fdbdatacube = gj.GribJump()
+        self.slicer = HullSlicer()
+        self.API = Polytope(
+            datacube=self.fdbdatacube,
+            engine=self.slicer,
+            options=self.options,
+        )
+        result = self.API.retrieve(request)
+        result.pprint()
+        assert len(result.leaves) == 8
+
+    @pytest.mark.fdb
+    def test_fdb_datacube_surrounding(self):
+        import pygribjump as gj
+
+        request = Request(
+            Select("step", [0]),
+            Select("levtype", ["sfc"]),
+            Span("date", pd.Timestamp("20230625T120000"), pd.Timestamp("20230626T120000")),
+            Select("domain", ["g"]),
+            Select("expver", ["0001"]),
+            Select("param", ["167"]),
+            Select("class", ["od"]),
+            Select("stream", ["oper"]),
+            Select("type", ["an"]),
+            Union(
+                ["latitude", "longitude"],
+                Point(["latitude", "longitude"], [[25, 30]], method="surrounding"),
+                Point(["latitude", "longitude"], [[-15, -45]], method="surrounding"),
+            ),
+        )
+
+        self.fdbdatacube = gj.GribJump()
+        self.slicer = HullSlicer()
+        self.API = Polytope(
+            datacube=self.fdbdatacube,
+            engine=self.slicer,
+            options=self.options,
+        )
+        result = self.API.retrieve(request)
+        result.pprint()
+        assert len(result.leaves) == 4
+        tot_leaves = 0
+        for leaf in result.leaves:
+            tot_leaves += len(leaf.result)
+        assert tot_leaves == 9
+
+    # @pytest.mark.fdb
+    # def test_fdb_datacube_mix_methods(self):
+    #     import pygribjump as gj
+
+    #     request = Request(
+    #         Select("step", [0]),
+    #         Select("levtype", ["sfc"]),
+    #         Span("date", pd.Timestamp("20230625T120000"), pd.Timestamp("20230626T120000")),
+    #         Select("domain", ["g"]),
+    #         Select("expver", ["0001"]),
+    #         Select("param", ["167"]),
+    #         Select("class", ["od"]),
+    #         Select("stream", ["oper"]),
+    #         Select("type", ["an"]),
+    #         Union(["latitude", "longitude"],
+    #               Point(["latitude", "longitude"], [[25, 30]], method="nearest"),
+    #               Point(["latitude", "longitude"], [[-15, -45]], method="surrounding"))
+    #     )
+
+    #     self.fdbdatacube = gj.GribJump()
+    #     self.slicer = HullSlicer()
+    #     self.API = Polytope(
+    #         datacube=self.fdbdatacube,
+    #         engine=self.slicer,
+    #         options=self.options,
+    #     )
+    #     result = self.API.retrieve(request)
+    #     result.pprint()
+    #     assert len(result.leaves) == 6