Feature/updated protobuf w develop

performance/fdb_slice_many_numbers_timeseries.py

@@ -14,7 +14,9 @@
     "step": {"type_change": "int"},
     "number": {"type_change": "int"},
     "longitude": {"cyclic": [0, 360]},
+    "latitude": {"reverse": {True}},
 }
+
 config = {"class": "od", "expver": "0001", "levtype": "sfc", "type": "pf"}
 fdbdatacube = FDBDatacube(config, axis_options=options)
 self_API = Polytope(datacube=fdbdatacube, axis_options=options)
@@ -39,7 +41,12 @@
     Point(["latitude", "longitude"], [[0.04, 0]], method="surrounding"),
     All("number"),
 )
+time3 = time.time()
 result = self_API.retrieve(request)
+time4 = time.time()
 print(time.time() - time1)
 print(time.time() - time2)
+print(time4 - time3)
 print(len(result.leaves))
+print([len(leaf.result) for leaf in result.leaves])
+result.pprint()

polytope/datacube/backends/datacube.py

@@ -1,30 +1,23 @@
-import argparse
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, List, Literal, Optional, Union
-
-import xarray as xr
-from conflator import ConfigModel, Conflator
-from pydantic import ConfigDict
+from typing import Any
 
 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,
 )
 
 
 class Datacube(ABC):
-    def __init__(self, axis_options=None, datacube_options=None):
+    def __init__(self, axis_options=None, compressed_axes_options=[]):
         if axis_options is None:
             self.axis_options = {}
         else:
             self.axis_options = axis_options
-        if datacube_options is None:
-            datacube_options = {}
-        self.axis_with_identical_structure_after = datacube_options.get("identical structure after")
         self.coupled_axes = []
         self.axis_counter = 0
         self.complete_axes = []
@@ -34,10 +27,13 @@ def __init__(self, axis_options=None, datacube_options=None):
         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
@@ -57,10 +53,27 @@ def _create_axes(self, name, values, transformation_type_key, transformation_opt
         transformation = DatacubeAxisTransformation.create_transform(
             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:
+                    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
@@ -132,57 +145,15 @@ def remap_path(self, path: DatacubePath):
         return path
 
     @staticmethod
-    def create_axes_config(axis_options):
-        class TransformationConfig(ConfigModel):
-            model_config = ConfigDict(extra="forbid")
-            name: str = ""
-
-        class CyclicConfig(TransformationConfig):
-            name: Literal["cyclic"]
-            range: List[float] = [0]
-
-        class MapperConfig(TransformationConfig):
-            name: Literal["mapper"]
-            type: str = ""
-            resolution: Union[int, List[int]] = 0
-            axes: List[str] = [""]
-            local: Optional[List[float]] = None
-
-        class ReverseConfig(TransformationConfig):
-            name: Literal["reverse"]
-            is_reverse: bool = False
-
-        class TypeChangeConfig(TransformationConfig):
-            name: Literal["type_change"]
-            type: str = "int"
-
-        class MergeConfig(TransformationConfig):
-            name: Literal["merge"]
-            other_axis: str = ""
-            linkers: List[str] = [""]
-
-        action_subclasses_union = Union[CyclicConfig, MapperConfig, ReverseConfig, TypeChangeConfig, MergeConfig]
-
-        class AxisConfig(ConfigModel):
-            axis_name: str = ""
-            transformations: list[action_subclasses_union]
-
-        class Config(ConfigModel):
-            config: list[AxisConfig] = []
-
-        parser = argparse.ArgumentParser(allow_abbrev=False)
-        axis_config = Conflator(app_name="polytope", model=Config, cli=False, argparser=parser).load()
-        if axis_options.get("config"):
-            axis_config = Config(config=axis_options.get("config"))
-
-        return axis_config
-
-    @staticmethod
-    def create(datacube, axis_options: dict, datacube_options={}):
-        if isinstance(datacube, (xr.core.dataarray.DataArray, xr.core.dataset.Dataset)):
+    def create(request, datacube, config={}, axis_options={}, compressed_axes_options=[]):
+        # 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, request, config, axis_options, compressed_axes_options)
+            return fdbdatacube

polytope/datacube/backends/fdb.py

@@ -1,27 +1,27 @@
 import logging
 from copy import deepcopy
-
-import pygribjump as pygj
+from itertools import product
 
 from ...utility.geometry import nearest_pt
-from .datacube import Datacube, IndexTree
+from .datacube import Datacube, TensorIndexTree
 
 
 class FDBDatacube(Datacube):
-    def __init__(self, request, config=None, axis_options=None, datacube_options=None):
+    def __init__(self, gj, request, config=None, axis_options=None, compressed_axes_options=[]):
         if config is None:
             config = {}
 
-        super().__init__(axis_options, datacube_options)
+        super().__init__(axis_options, compressed_axes_options)
 
         logging.info("Created an FDB datacube with options: " + str(axis_options))
 
         self.unwanted_path = {}
-        self.axis_options = Datacube.create_axes_config(axis_options).config
+        self.axis_options = axis_options
 
         partial_request = config
         # Find values in the level 3 FDB datacube
-        self.gj = pygj.GribJump()
+
+        self.gj = gj
         self.fdb_coordinates = self.gj.axes(partial_request)
 
         self.check_branching_axes(request)
@@ -61,20 +61,45 @@ def check_branching_axes(self, request):
                     (upper, lower, idx) = polytope.extents(ax)
                     if "sfc" in polytope.points[idx]:
                         self.fdb_coordinates.pop("levelist")
-        print(self.fdb_coordinates)
         self.fdb_coordinates.pop("quantile", None)
-        print(self.fdb_coordinates)
 
-    def get(self, requests: IndexTree):
+    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)
-        output_values = self.gj.extract(fdb_requests)
-        self.assign_fdb_output_to_nodes(output_values, fdb_requests_decoding_info)
 
-    def get_fdb_requests(self, requests: IndexTree, fdb_requests=[], fdb_requests_decoding_info=[], leaf_path=None):
+        # TODO: note that this doesn't exactly work as intended, it's just going to retrieve value from gribjump that
+        # corresponds to first value in the compressed tuples
+
+        # TODO: here, loop through the fdb requests and request from gj and directly add to the nodes
+
+        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():
+                # remove the tuple of the request when we ask the fdb
+                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])
+                output_values = self.gj.extract([complete_uncompressed_request])
+                self.assign_fdb_output_to_nodes(output_values, [fdb_requests_decoding_info[j]])
+
+    def get_fdb_requests(
+        self, requests: TensorIndexTree, fdb_requests=[], fdb_requests_decoding_info=[], leaf_path=None
+    ):
         if leaf_path is None:
             leaf_path = {}
 
@@ -86,7 +111,7 @@ def get_fdb_requests(self, requests: IndexTree, fdb_requests=[], fdb_requests_de
                 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.value}
+            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
@@ -136,7 +161,7 @@ def get_2nd_last_values(self, requests, leaf_path=None):
             found_latlon_pts = []
             for lat_child in requests.children:
                 for lon_child in lat_child.children:
-                    found_latlon_pts.append([lat_child.value, lon_child.value])
+                    found_latlon_pts.append([lat_child.values, lon_child.values])
 
             # now find the nearest lat lon to the points requested
             nearest_latlons = []
@@ -145,20 +170,21 @@ def get_2nd_last_values(self, requests, leaf_path=None):
                 nearest_latlons.append(nearest_latlon)
 
             # need to remove the branches that do not fit
-            lat_children_values = [child.value for child in requests.children]
+            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.value == lat_child_val][0]
-                if lat_child.value not in [latlon[0] for latlon in nearest_latlons]:
+                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.value]
-                    lon_children_values = [child.value for child in lat_child.children]
+                    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.value == lon_child_val][0]
-                        if lon_child.value not in possible_lons:
-                            lon_child.remove_branch()
+                        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)
 
         lat_length = len(requests.children)
         range_lengths = [False] * lat_length
@@ -167,13 +193,13 @@ def get_2nd_last_values(self, requests, leaf_path=None):
         for i in range(len(requests.children)):
             lat_child = requests.children[i]
             lon_length = len(lat_child.children)
-            range_lengths[i] = [1] * lon_length
+            range_lengths[i] = [0] * lon_length
             current_start_idxs[i] = [None] * lon_length
-            fdb_node_ranges[i] = [[IndexTree.root] * lon_length] * lon_length
+            fdb_node_ranges[i] = [[TensorIndexTree.root for y in range(lon_length)] for x in range(lon_length)]
             range_length = deepcopy(range_lengths[i])
             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.value}
+            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
@@ -184,30 +210,31 @@ def get_2nd_last_values(self, requests, leaf_path=None):
             )
 
         leaf_path_copy = deepcopy(leaf_path)
-        leaf_path_copy.pop("values")
-        # leaf_path_copy.update(self.necessary_popped_axes)
+        leaf_path_copy.pop("values", None)
         return (leaf_path_copy, range_lengths, current_start_idxs, fdb_node_ranges, lat_length)
 
     def get_last_layer_before_leaf(self, requests, leaf_path, range_l, current_idx, fdb_range_n):
         i = 0
         for c in requests.children:
+            fdb_range_n_i = fdb_range_n[i]
             # now c are the leaves of the initial tree
-            key_value_path = {c.axis.name: c.value}
+            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
             )
             leaf_path.update(key_value_path)
             last_idx = key_value_path["values"]
             if current_idx[i] is None:
+                range_l[i] = 1
                 current_idx[i] = last_idx
-                fdb_range_n[i][range_l[i] - 1] = c
+                fdb_range_n_i[range_l[i] - 1] = c
             else:
                 if last_idx == current_idx[i] + range_l[i]:
                     range_l[i] += 1
-                    fdb_range_n[i][range_l[i] - 1] = c
+                    fdb_range_n_i[range_l[i] - 1] = c
                 else:
-                    key_value_path = {c.axis.name: c.value}
+                    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
@@ -216,6 +243,8 @@ def get_last_layer_before_leaf(self, requests, leaf_path, range_l, current_idx,
                     i += 1
                     current_start_idx = key_value_path["values"]
                     current_idx[i] = current_start_idx
+                    range_l[i] = 1
+                    fdb_range_n[i][range_l[i] - 1] = c
         return (range_l, current_idx, fdb_range_n)
 
     def assign_fdb_output_to_nodes(self, output_values, fdb_requests_decoding_info):
@@ -240,7 +269,7 @@ def assign_fdb_output_to_nodes(self, output_values, fdb_requests_decoding_info):
             for i in range(len(sorted_fdb_range_nodes)):
                 for j in range(sorted_range_lengths[i]):
                     n = sorted_fdb_range_nodes[i][j]
-                    n.result = request_output_values[0][i][0][j]
+                    n.result.append(request_output_values[0][i][0][j])
 
     def sort_fdb_request_ranges(self, range_lengths, current_start_idx, lat_length):
         interm_request_ranges = []
@@ -263,3 +292,21 @@ def select(self, path, unmapped_path):
 
     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