polytope-python 1.0.3__tar.gz → 1.0.5__tar.gz

{polytope-python-1.0.3/polytope_python.egg-info → polytope-python-1.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.0
 Name: polytope-python
-Version: 1.0.3
+Version: 1.0.5
 Summary: Polytope datacube feature extraction library
 Home-page: https://github.com/ecmwf/polytope
 Author: ECMWF

{polytope-python-1.0.3 → polytope-python-1.0.5}/polytope/datacube/backends/datacube.py

@@ -1,14 +1,11 @@
-import importlib
 import logging
-import math
 from abc import ABC, abstractmethod
 from typing import Any
 
-import xarray as xr
-
-from ...utility.combinatorics import unique, validate_axes
+from ...utility.combinatorics import validate_axes
 from ..datacube_axis import DatacubeAxis
-from ..index_tree import DatacubePath, IndexTree
+from ..tensor_index_tree import DatacubePath, TensorIndexTree
+from ..transformations.datacube_mappers.datacube_mappers import DatacubeMapper
 from ..transformations.datacube_transformations import (
     DatacubeAxisTransformation,
     has_transform,
@@ -16,8 +13,27 @@ from ..transformations.datacube_transformations import (
 
 
 class Datacube(ABC):
+    def __init__(self, axis_options=None, compressed_axes_options=[]):
+        if axis_options is None:
+            self.axis_options = {}
+        else:
+            self.axis_options = axis_options
+        self.coupled_axes = []
+        self.axis_counter = 0
+        self.complete_axes = []
+        self.blocked_axes = []
+        self.fake_axes = []
+        self.treated_axes = []
+        self.nearest_search = {}
+        self._axes = None
+        self.transformed_axes = []
+        self.compressed_grid_axes = []
+        self.merged_axes = []
+        self.unwanted_path = {}
+        self.compressed_axes = compressed_axes_options
+
     @abstractmethod
-    def get(self, requests: IndexTree) -> Any:
+    def get(self, requests: TensorIndexTree) -> Any:
         """Return data given a set of request trees"""
 
     @property
@@ -30,39 +46,55 @@ class Datacube(ABC):
 
     def _create_axes(self, name, values, transformation_type_key, transformation_options):
         # first check what the final axes are for this axis name given transformations
+        transformation_options = transformation_type_key
         final_axis_names = DatacubeAxisTransformation.get_final_axes(
-            name, transformation_type_key, transformation_options
+            name, transformation_type_key.name, transformation_options
         )
         transformation = DatacubeAxisTransformation.create_transform(
-            name, transformation_type_key, transformation_options
+            name, transformation_type_key.name, transformation_options
         )
+
+        # do not compress merged axes
+        if transformation_type_key.name == "merge":
+            self.merged_axes.append(name)
+            self.merged_axes.append(final_axis_names)
+            for axis in final_axis_names:
+                # remove the merged_axes from the possible compressed axes
+                if axis in self.compressed_axes:
+                    self.compressed_axes.remove(axis)
+
         for blocked_axis in transformation.blocked_axes():
             self.blocked_axes.append(blocked_axis)
+        if isinstance(transformation, DatacubeMapper):
+            # TODO: do we use this?? This shouldn't work for a disk in lat/lon on a octahedral or other grid??
+            for compressed_grid_axis in transformation.compressed_grid_axes:
+                self.compressed_grid_axes.append(compressed_grid_axis)
         if len(final_axis_names) > 1:
             self.coupled_axes.append(final_axis_names)
+            for axis in final_axis_names:
+                if axis in self.compressed_axes and axis != final_axis_names[-1]:
+                    self.compressed_axes.remove(axis)
         for axis_name in final_axis_names:
             self.fake_axes.append(axis_name)
             # if axis does not yet exist, create it
 
             # first need to change the values so that we have right type
             values = transformation.change_val_type(axis_name, values)
-            if self._axes is None:
-                DatacubeAxis.create_standard(axis_name, values, self)
-            elif axis_name not in self._axes.keys():
+            if self._axes is None or axis_name not in self._axes.keys():
                 DatacubeAxis.create_standard(axis_name, values, self)
             # add transformation tag to axis, as well as transformation options for later
-            setattr(self._axes[axis_name], has_transform[transformation_type_key], True)  # where has_transform is a
-            # factory inside datacube_transformations to set the has_transform, is_cyclic etc axis properties
+            setattr(self._axes[axis_name], has_transform[transformation_type_key.name], True)  # where has_transform is
+            # a factory inside datacube_transformations to set the has_transform, is_cyclic etc axis properties
             # add the specific transformation handled here to the relevant axes
             # Modify the axis to update with the tag
-            decorator_module = importlib.import_module("polytope.datacube.datacube_axis")
-            decorator = getattr(decorator_module, transformation_type_key)
-            decorator(self._axes[axis_name])
+
             if transformation not in self._axes[axis_name].transformations:  # Avoids duplicates being stored
                 self._axes[axis_name].transformations.append(transformation)
 
     def _add_all_transformation_axes(self, options, name, values):
-        for transformation_type_key in options.keys():
+        for transformation_type_key in options.transformations:
+            if transformation_type_key != "cyclic":
+                self.transformed_axes.append(name)
             self._create_axes(name, values, transformation_type_key, options)
 
     def _check_and_add_axes(self, options, name, values):
@@ -70,9 +102,7 @@ class Datacube(ABC):
             self._add_all_transformation_axes(options, name, values)
         else:
             if name not in self.blocked_axes:
-                if self._axes is None:
-                    DatacubeAxis.create_standard(name, values, self)
-                elif name not in self._axes.keys():
+                if self._axes is None or name not in self._axes.keys():
                     DatacubeAxis.create_standard(name, values, self)
 
     def has_index(self, path: DatacubePath, axis, index):
@@ -96,46 +126,12 @@ class Datacube(ABC):
         """
         path = self.fit_path(path)
         indexes = axis.find_indexes(path, self)
-        search_ranges = axis.remap([lower, upper])
-        original_search_ranges = axis.to_intervals([lower, upper])
-        # Find the offsets for each interval in the requested range, which we will need later
-        search_ranges_offset = []
-        for r in original_search_ranges:
-            offset = axis.offset(r)
-            search_ranges_offset.append(offset)
-        idx_between = self._look_up_datacube(search_ranges, search_ranges_offset, indexes, axis, method)
-        # Remove duplicates even if difference of the order of the axis tolerance
-        if offset is not None:
-            # Note that we can only do unique if not dealing with time values
-            idx_between = unique(idx_between)
+        idx_between = axis.find_indices_between(indexes, lower, upper, self, method)
 
         logging.info(f"For axis {axis.name} between {lower} and {upper}, found indices {idx_between}")
 
         return idx_between
 
-    def _look_up_datacube(self, search_ranges, search_ranges_offset, indexes, axis, method):
-        idx_between = []
-        for i in range(len(search_ranges)):
-            r = search_ranges[i]
-            offset = search_ranges_offset[i]
-            low = r[0]
-            up = r[1]
-            indexes_between = axis.find_indices_between([indexes], low, up, self, method)
-            # Now the indexes_between are values on the cyclic range so need to remap them to their original
-            # values before returning them
-            for j in range(len(indexes_between)):
-                # if we have a special indexes between range that needs additional offset, treat it here
-                if len(indexes_between[j]) == 0:
-                    idx_between = idx_between
-                else:
-                    for k in range(len(indexes_between[j])):
-                        if offset is None:
-                            indexes_between[j][k] = indexes_between[j][k]
-                        else:
-                            indexes_between[j][k] = round(indexes_between[j][k] + offset, int(-math.log10(axis.tol)))
-                        idx_between.append(indexes_between[j][k])
-        return idx_between
-
     def get_mapper(self, axis):
         """
         Get the type mapper for a subaxis of the datacube given by label
@@ -149,11 +145,18 @@ class Datacube(ABC):
         return path
 
     @staticmethod
-    def create(datacube, axis_options: dict, datacube_options={}):
-        if isinstance(datacube, (xr.core.dataarray.DataArray, xr.core.dataset.Dataset)):
+    def create(datacube, config={}, axis_options={}, compressed_axes_options=[], alternative_axes=[]):
+        # TODO: get the configs as None for pre-determined value and change them to empty dictionary inside the function
+        if type(datacube).__name__ == "DataArray":
             from .xarray import XArrayDatacube
 
-            xadatacube = XArrayDatacube(datacube, axis_options, datacube_options)
+            xadatacube = XArrayDatacube(datacube, axis_options, compressed_axes_options)
             return xadatacube
-        else:
-            return datacube
+        if type(datacube).__name__ == "GribJump":
+            from .fdb import FDBDatacube
+
+            fdbdatacube = FDBDatacube(datacube, config, axis_options, compressed_axes_options, alternative_axes)
+            return fdbdatacube
+
+    def check_branching_axes(self, request):
+        pass

polytope-python-1.0.5/polytope/datacube/backends/fdb.py

@@ -0,0 +1,357 @@
+import logging
+import operator
+from copy import deepcopy
+from itertools import product
+
+from ...utility.geometry import nearest_pt
+from .datacube import Datacube, TensorIndexTree
+
+
+class FDBDatacube(Datacube):
+    def __init__(self, gj, config=None, axis_options=None, compressed_axes_options=[], alternative_axes=[]):
+        if config is None:
+            config = {}
+
+        super().__init__(axis_options, compressed_axes_options)
+
+        logging.info("Created an FDB datacube with options: " + str(axis_options))
+
+        self.unwanted_path = {}
+        self.axis_options = axis_options
+
+        partial_request = config
+        # Find values in the level 3 FDB datacube
+
+        self.gj = gj
+        if len(alternative_axes) == 0:
+            self.fdb_coordinates = self.gj.axes(partial_request)
+        else:
+            self.fdb_coordinates = {}
+            for axis_config in alternative_axes:
+                self.fdb_coordinates[axis_config.axis_name] = axis_config.values
+
+        logging.info("Axes returned from GribJump are: " + str(self.fdb_coordinates))
+
+        self.fdb_coordinates["values"] = []
+        for name, values in self.fdb_coordinates.items():
+            values.sort()
+            options = None
+            for opt in self.axis_options:
+                if opt.axis_name == name:
+                    options = opt
+
+            self._check_and_add_axes(options, name, values)
+            self.treated_axes.append(name)
+            self.complete_axes.append(name)
+
+        # add other options to axis which were just created above like "lat" for the mapper transformations for eg
+        for name in self._axes:
+            if name not in self.treated_axes:
+                options = None
+                for opt in self.axis_options:
+                    if opt.axis_name == name:
+                        options = opt
+
+                val = self._axes[name].type
+                self._check_and_add_axes(options, name, val)
+
+        logging.info("Polytope created axes for %s", self._axes.keys())
+
+    def check_branching_axes(self, request):
+        polytopes = request.polytopes()
+        for polytope in polytopes:
+            for ax in polytope._axes:
+                if ax == "levtype":
+                    (upper, lower, idx) = polytope.extents(ax)
+                    if "sfc" in polytope.points[idx]:
+                        self.fdb_coordinates.pop("levelist", None)
+        self.fdb_coordinates.pop("quantile", None)
+        # TODO: When do these not appear??
+        self.fdb_coordinates.pop("direction", None)
+        self.fdb_coordinates.pop("frequency", None)
+
+        # NOTE: verify that we also remove the axis object for axes we've removed here
+        axes_to_remove = set(self.complete_axes) - set(self.fdb_coordinates.keys())
+
+        # Remove the keys from self._axes
+        for axis_name in axes_to_remove:
+            self._axes.pop(axis_name, None)
+
+    def get(self, requests: TensorIndexTree):
+        if len(requests.children) == 0:
+            return requests
+        fdb_requests = []
+        fdb_requests_decoding_info = []
+        self.get_fdb_requests(requests, fdb_requests, fdb_requests_decoding_info)
+
+        # here, loop through the fdb requests and request from gj and directly add to the nodes
+        complete_list_complete_uncompressed_requests = []
+        complete_fdb_decoding_info = []
+        for j, compressed_request in enumerate(fdb_requests):
+            uncompressed_request = {}
+
+            # Need to determine the possible decompressed requests
+
+            # find the possible combinations of compressed indices
+            interm_branch_tuple_values = []
+            for key in compressed_request[0].keys():
+                interm_branch_tuple_values.append(compressed_request[0][key])
+            request_combis = product(*interm_branch_tuple_values)
+
+            # Need to extract the possible requests and add them to the right nodes
+            for combi in request_combis:
+                uncompressed_request = {}
+                for i, key in enumerate(compressed_request[0].keys()):
+                    uncompressed_request[key] = combi[i]
+                complete_uncompressed_request = (uncompressed_request, compressed_request[1])
+                complete_list_complete_uncompressed_requests.append(complete_uncompressed_request)
+                complete_fdb_decoding_info.append(fdb_requests_decoding_info[j])
+        logging.debug("The requests we give GribJump are: %s", complete_list_complete_uncompressed_requests)
+        output_values = self.gj.extract(complete_list_complete_uncompressed_requests)
+        logging.debug("GribJump outputs: %s", output_values)
+        self.assign_fdb_output_to_nodes(output_values, complete_fdb_decoding_info)
+
+    def get_fdb_requests(
+        self,
+        requests: TensorIndexTree,
+        fdb_requests=[],
+        fdb_requests_decoding_info=[],
+        leaf_path=None,
+    ):
+        if leaf_path is None:
+            leaf_path = {}
+
+        # First when request node is root, go to its children
+        if requests.axis.name == "root":
+            logging.info("Looking for data for the tree: %s", [leaf.flatten() for leaf in requests.leaves])
+
+            for c in requests.children:
+                self.get_fdb_requests(c, fdb_requests, fdb_requests_decoding_info)
+        # If request node has no children, we have a leaf so need to assign fdb values to it
+        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].children[0].children) == 0:
+                # find the fdb_requests and associated nodes to which to add results
+                (path, current_start_idxs, fdb_node_ranges, lat_length) = self.get_2nd_last_values(requests, leaf_path)
+                (
+                    original_indices,
+                    sorted_request_ranges,
+                    fdb_node_ranges,
+                ) = self.sort_fdb_request_ranges(current_start_idxs, lat_length, fdb_node_ranges)
+                fdb_requests.append((path, sorted_request_ranges))
+                fdb_requests_decoding_info.append((original_indices, fdb_node_ranges))
+
+            # Otherwise remap the path for this key and iterate again over children
+            else:
+                for c in requests.children:
+                    self.get_fdb_requests(c, fdb_requests, fdb_requests_decoding_info, leaf_path)
+
+    def remove_duplicates_in_request_ranges(self, fdb_node_ranges, current_start_idxs):
+        seen_indices = set()
+        for i, idxs_list in enumerate(current_start_idxs):
+            for k, sub_lat_idxs in enumerate(idxs_list):
+                actual_fdb_node = fdb_node_ranges[i][k]
+                original_fdb_node_range_vals = []
+                new_current_start_idx = []
+                for j, idx in enumerate(sub_lat_idxs):
+                    if idx not in seen_indices:
+                        # TODO: need to remove it from the values in the corresponding tree node
+                        # TODO: need to read just the range we give to gj ... DONE?
+                        original_fdb_node_range_vals.append(actual_fdb_node[0].values[j])
+                        seen_indices.add(idx)
+                        new_current_start_idx.append(idx)
+                if original_fdb_node_range_vals != []:
+                    actual_fdb_node[0].values = tuple(original_fdb_node_range_vals)
+                else:
+                    # there are no values on this node anymore so can remove it
+                    actual_fdb_node[0].remove_branch()
+                if len(new_current_start_idx) == 0:
+                    current_start_idxs[i].pop(k)
+                else:
+                    current_start_idxs[i][k] = new_current_start_idx
+        return (fdb_node_ranges, current_start_idxs)
+
+    def nearest_lat_lon_search(self, requests):
+        if len(self.nearest_search) != 0:
+            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():
+                raise Exception("nearest point search axes are wrong")
+
+            second_ax = requests.children[0].children[0].axis
+
+            # TODO: actually, here we should not remap the nearest_pts, we should instead unmap the
+            # found_latlon_pts and then remap them later once we have compared found_latlon_pts and nearest_pts
+            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]
+                )
+            ]
+
+            found_latlon_pts = []
+            for lat_child in requests.children:
+                for lon_child in lat_child.children:
+                    found_latlon_pts.append([lat_child.values, lon_child.values])
+
+            # now find the nearest lat lon to the points requested
+            nearest_latlons = []
+            for pt in nearest_pts:
+                nearest_latlon = nearest_pt(found_latlon_pts, pt)
+                nearest_latlons.append(nearest_latlon)
+
+            # need to remove the branches that do not fit
+            lat_children_values = [child.values for child in requests.children]
+            for i in range(len(lat_children_values)):
+                lat_child_val = lat_children_values[i]
+                lat_child = [child for child in requests.children if child.values == lat_child_val][0]
+                if lat_child.values not in [(latlon[0],) for latlon in nearest_latlons]:
+                    lat_child.remove_branch()
+                else:
+                    possible_lons = [latlon[1] for latlon in nearest_latlons if (latlon[0],) == lat_child.values]
+                    lon_children_values = [child.values for child in lat_child.children]
+                    for j in range(len(lon_children_values)):
+                        lon_child_val = lon_children_values[j]
+                        lon_child = [child for child in lat_child.children if child.values == lon_child_val][0]
+                        for value in lon_child.values:
+                            if value not in possible_lons:
+                                lon_child.remove_compressed_branch(value)
+
+    def get_2nd_last_values(self, requests, leaf_path=None):
+        if leaf_path is None:
+            leaf_path = {}
+        # In this function, we recursively loop over the last two layers of the tree and store the indices of the
+        # request ranges in those layers
+        self.nearest_lat_lon_search(requests)
+
+        lat_length = len(requests.children)
+        current_start_idxs = [False] * lat_length
+        fdb_node_ranges = [False] * lat_length
+        for i in range(len(requests.children)):
+            lat_child = requests.children[i]
+            lon_length = len(lat_child.children)
+            current_start_idxs[i] = [None] * lon_length
+            fdb_node_ranges[i] = [[TensorIndexTree.root for y in range(lon_length)] for x in range(lon_length)]
+            current_start_idx = deepcopy(current_start_idxs[i])
+            fdb_range_nodes = deepcopy(fdb_node_ranges[i])
+            key_value_path = {lat_child.axis.name: lat_child.values}
+            ax = lat_child.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)
+            (current_start_idxs[i], fdb_node_ranges[i]) = self.get_last_layer_before_leaf(
+                lat_child, leaf_path, current_start_idx, fdb_range_nodes
+            )
+
+        leaf_path_copy = deepcopy(leaf_path)
+        leaf_path_copy.pop("values", None)
+        return (leaf_path_copy, current_start_idxs, fdb_node_ranges, lat_length)
+
+    def get_last_layer_before_leaf(self, requests, leaf_path, current_idx, fdb_range_n):
+        current_idx = [[] for i in range(len(requests.children))]
+        fdb_range_n = [[] for i in range(len(requests.children))]
+        for i, c in enumerate(requests.children):
+            # now c are the leaves of the initial tree
+            key_value_path = {c.axis.name: c.values}
+            ax = c.axis
+            (key_value_path, leaf_path, self.unwanted_path) = ax.unmap_path_key(
+                key_value_path, leaf_path, self.unwanted_path
+            )
+            # TODO: change this to accommodate non consecutive indexes being compressed too
+            current_idx[i].extend(key_value_path["values"])
+            fdb_range_n[i].append(c)
+        return (current_idx, fdb_range_n)
+
+    def assign_fdb_output_to_nodes(self, output_values, fdb_requests_decoding_info):
+        for k in range(len(output_values)):
+            request_output_values = output_values[k]
+            (
+                original_indices,
+                fdb_node_ranges,
+            ) = fdb_requests_decoding_info[k]
+            sorted_fdb_range_nodes = [fdb_node_ranges[i] for i in original_indices]
+            for i in range(len(sorted_fdb_range_nodes)):
+                n = sorted_fdb_range_nodes[i][0]
+                if len(request_output_values[0]) == 0:
+                    # If we are here, no data was found for this path in the fdb
+                    none_array = [None] * len(n.values)
+                    n.result.extend(none_array)
+                else:
+                    interm_request_output_values = request_output_values[0][i][0]
+                    n.result.extend(interm_request_output_values)
+
+    def sort_fdb_request_ranges(self, current_start_idx, lat_length, fdb_node_ranges):
+        (new_fdb_node_ranges, new_current_start_idx) = self.remove_duplicates_in_request_ranges(
+            fdb_node_ranges, current_start_idx
+        )
+        interm_request_ranges = []
+        # TODO: modify the start indexes to have as many arrays as the request ranges
+        new_fdb_node_ranges = []
+        for i in range(lat_length):
+            interm_fdb_nodes = fdb_node_ranges[i]
+            old_interm_start_idx = current_start_idx[i]
+            for j in range(len(old_interm_start_idx)):
+                # TODO: if we sorted the cyclic values in increasing order on the tree too,
+                # then we wouldn't have to sort here?
+                sorted_list = sorted(enumerate(old_interm_start_idx[j]), key=lambda x: x[1])
+                original_indices_idx, interm_start_idx = zip(*sorted_list)
+                for interm_fdb_nodes_obj in interm_fdb_nodes[j]:
+                    interm_fdb_nodes_obj.values = tuple([interm_fdb_nodes_obj.values[k] for k in original_indices_idx])
+                if abs(interm_start_idx[-1] + 1 - interm_start_idx[0]) <= len(interm_start_idx):
+                    current_request_ranges = (interm_start_idx[0], interm_start_idx[-1] + 1)
+                    interm_request_ranges.append(current_request_ranges)
+                    new_fdb_node_ranges.append(interm_fdb_nodes[j])
+                else:
+                    jumps = list(map(operator.sub, interm_start_idx[1:], interm_start_idx[:-1]))
+                    last_idx = 0
+                    for k, jump in enumerate(jumps):
+                        if jump > 1:
+                            current_request_ranges = (interm_start_idx[last_idx], interm_start_idx[k] + 1)
+                            new_fdb_node_ranges.append(interm_fdb_nodes[j])
+                            last_idx = k + 1
+                            interm_request_ranges.append(current_request_ranges)
+                        if k == len(interm_start_idx) - 2:
+                            current_request_ranges = (interm_start_idx[last_idx], interm_start_idx[-1] + 1)
+                            interm_request_ranges.append(current_request_ranges)
+                            new_fdb_node_ranges.append(interm_fdb_nodes[j])
+        request_ranges_with_idx = list(enumerate(interm_request_ranges))
+        sorted_list = sorted(request_ranges_with_idx, key=lambda x: x[1][0])
+        original_indices, sorted_request_ranges = zip(*sorted_list)
+        logging.debug("We sorted the request ranges into: %s", sorted_request_ranges)
+        logging.debug("The sorted and unique leaf node ranges are: %s", new_fdb_node_ranges)
+        return (original_indices, sorted_request_ranges, new_fdb_node_ranges)
+
+    def datacube_natural_indexes(self, axis, subarray):
+        indexes = subarray.get(axis.name, None)
+        return indexes
+
+    def select(self, path, unmapped_path):
+        return self.fdb_coordinates
+
+    def ax_vals(self, name):
+        return self.fdb_coordinates.get(name, None)
+
+    def prep_tree_encoding(self, node, unwanted_path=None):
+        # TODO: prepare the tree for protobuf encoding
+        # ie transform all axes for gribjump and adding the index property on the leaves
+        if unwanted_path is None:
+            unwanted_path = {}
+
+        ax = node.axis
+        (new_node, unwanted_path) = ax.unmap_tree_node(node, unwanted_path)
+
+        if len(node.children) != 0:
+            for c in new_node.children:
+                self.prep_tree_encoding(c, unwanted_path)
+
+    def prep_tree_decoding(self, tree):
+        # TODO: transform the tree after decoding from protobuf
+        # ie unstransform all axes from gribjump and put the indexes back as a leaf/extra node
+        pass

{polytope-python-1.0.3 → polytope-python-1.0.5}/polytope/datacube/backends/mock.py

@@ -3,11 +3,12 @@ from copy import deepcopy
 
 from ...utility.combinatorics import validate_axes
 from ..datacube_axis import IntDatacubeAxis
-from .datacube import Datacube, DatacubePath, IndexTree
+from .datacube import Datacube, DatacubePath, TensorIndexTree
 
 
 class MockDatacube(Datacube):
-    def __init__(self, dimensions, datacube_options={}):
+    def __init__(self, dimensions, compressed_axes_options=[]):
+        super().__init__({}, compressed_axes_options)
         assert isinstance(dimensions, dict)
 
         self.dimensions = dimensions
@@ -22,17 +23,15 @@ class MockDatacube(Datacube):
         for k, v in reversed(dimensions.items()):
             self.stride[k] = stride_cumulative
             stride_cumulative *= self.dimensions[k]
-        self.coupled_axes = []
-        self.axis_with_identical_structure_after = ""
 
-    def get(self, requests: IndexTree):
+    def get(self, requests: TensorIndexTree):
         # Takes in a datacube and verifies the leaves of the tree are complete
         # (ie it found values for all datacube axis)
 
         for r in requests.leaves:
             path = r.flatten()
             if len(path.items()) == len(self.dimensions.items()):
-                result = 0
+                result = (0,)
                 for k, v in path.items():
                     result += v * self.stride[k]