merge the improved xarray backend with the compressed tree slicing

performance/performance_many_num_steps.py

@@ -62,4 +62,4 @@
 print(time4 - time3)
 print(len(result.leaves))
 print([len(leaf.result) for leaf in result.leaves])
-result.pprint()
+result.pprint()

polytope/datacube/backends/datacube.py

@@ -34,6 +34,7 @@ def __init__(self, axis_options=None, datacube_options=None):
         self.transformed_axes = []
         self.compressed_grid_axes = []
         self.compressed_axes = []
+        self.unwanted_path = {}
 
     @abstractmethod
     def get(self, requests: IndexTree) -> Any:

polytope/datacube/backends/fdb.py

@@ -18,7 +18,7 @@ def __init__(self, config=None, axis_options=None, datacube_options=None):
 
         logging.info("Created an FDB datacube with options: " + str(axis_options))
 
-        self.unwanted_path = {}
+        # self.unwanted_path = {}
 
         partial_request = config
         # Find values in the level 3 FDB datacube

polytope/datacube/backends/xarray.py

@@ -37,34 +37,42 @@ def __init__(self, dataarray: xr.DataArray, axis_options=None, datacube_options=
                 val = self._axes[name].type
                 self._check_and_add_axes(options, name, val)
 
-    def get(self, requests: IndexTree):
-        for r in requests.leaves:
-            path = r.flatten()
-            if len(path.items()) == self.axis_counter:
-                # TODO: need to undo the tuples in the path into actual paths with a single value that xarray can read
-                unmapped_path = {}
-                path_copy = deepcopy(path)
-                for key in path_copy:
-                    axis = self._axes[key]
-                    key_value_path = {key: path_copy[key]}
-                    (key_value_path, path, unmapped_path) = axis.unmap_path_key(key_value_path, path, unmapped_path)
-                path.update(key_value_path)
-
-                unmapped_path = {}
-                self.refit_path(path, unmapped_path, path)
-                for key in path:
-                    path[key] = list(path[key])
-                for key in unmapped_path:
-                    if isinstance(unmapped_path[key], tuple):
-                        unmapped_path[key] = list(unmapped_path[key])
-
-                subxarray = self.dataarray.sel(path, method="nearest")
-                subxarray = subxarray.sel(unmapped_path)
-                value = subxarray.values
-                key = subxarray.name
-                r.result = (key, value)
+    def get(self, requests, leaf_path=None, axis_counter=0):
+        if leaf_path is None:
+            leaf_path = {}
+        if requests.axis.name == "root":
+            for c in requests.children:
+                self.get(c, leaf_path, axis_counter+1)
+        else:
+            key_value_path = {requests.axis.name: requests.values}
+            ax = requests.axis
+            (key_value_path, leaf_path, self.unwanted_path) = ax.unmap_path_key(
+                key_value_path, leaf_path, self.unwanted_path
+            )
+            leaf_path.update(key_value_path)
+            if len(requests.children) != 0:
+                # We are not a leaf and we loop over
+                for c in requests.children:
+                    self.get(c, leaf_path, axis_counter+1)
             else:
-                r.remove_branch()
+                if self.axis_counter != axis_counter:
+                    requests.remove_branch()
+                else:
+                    # We are at a leaf and need to assign value to it
+                    leaf_path_copy = deepcopy(leaf_path)
+                    unmapped_path = {}
+                    self.refit_path(leaf_path_copy, unmapped_path, leaf_path)
+                    for key in leaf_path_copy:
+                        leaf_path_copy[key] = list(leaf_path_copy[key])
+                    for key in unmapped_path:
+                        if isinstance(unmapped_path[key], tuple):
+                            unmapped_path[key] = list(unmapped_path[key])
+                    subxarray = self.dataarray.sel(leaf_path_copy, method="nearest")
+                    subxarray = subxarray.sel(unmapped_path)
+                    # TODO: might be faster to just return the sub-xarray here?
+                    value = subxarray.values
+                    key = subxarray.name
+                    requests.result = (key, value)
 
     def datacube_natural_indexes(self, axis, subarray):
         if axis.name in self.complete_axes:

polytope/engine/hullslicer.py

@@ -86,9 +86,6 @@ def _build_sliceable_child(self, polytope, ax, node, datacube, lower, upper, nex
             if flattened.get(datacube.coupled_axes[0][0], None) is not None:
                 flattened_tuple = (datacube.coupled_axes[0][0], flattened.get(datacube.coupled_axes[0][0], None))
                 flattened = {flattened_tuple[0]: flattened_tuple[1]}
-            else:
-                # flattened = {}
-                pass
 
         values = self.axis_values_between.get((flattened_tuple, ax.name, lower, upper, method), None)
         if self.axis_values_between.get((flattened_tuple, ax.name, lower, upper, method), None) is None:

tests/test_healpix_mapper.py

@@ -17,6 +17,7 @@ def setup_method(self, method):
         self.options = {
             "values": {"mapper": {"type": "healpix", "resolution": 32, "axes": ["latitude", "longitude"]}},
             "longitude": {"cyclic": [0, 360]},
+            "latitude": {"reverse": True}
         }
         self.slicer = HullSlicer()
         self.API = Polytope(datacube=self.latlon_array, engine=self.slicer, axis_options=self.options)
@@ -31,7 +32,7 @@ def test_healpix_grid(self):
             Select("valid_time", ["2022-12-14T13:00:00"]),
         )
         result = self.API.retrieve(request)
-        result.pprint()
+        # result.pprint()
         assert len(result.leaves) == 40
 
         lats = []