Feature/irregular grid merged develop

.github/workflows/ci.yaml

@@ -155,8 +155,8 @@ jobs:
           pip install setuptools wheel twine
       - name: Build and publish
         env:
-          TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
-          TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+          TWINE_USERNAME: "__token__"
+          TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
         run: |
           python setup.py sdist
           twine upload dist/*

docs/requirements.txt

@@ -1,3 +1,3 @@
-jinja2<3.1.0
+jinja2>=3.1.3
 Markdown<3.2
 mkdocs>=1.0

examples/requirements_examples.txt

@@ -1,13 +1,13 @@
 -r ../requirements.txt
 -r ../tests/requirements_test.txt
 
-matplotlib==3.6.2
-matplotlib-inline==0.1.6
-Pillow==9.3.0
-Shapely==1.8.5.post1
-shp==1.0.2
-Fiona==1.8.22
-geopandas==0.12.2
-plotly==5.11.0
-pyshp==2.3.1
-cfgrib==0.9.10.3
+matplotlib
+matplotlib-inline
+Pillow
+Shapely
+shp
+Fiona
+geopandas
+plotly
+pyshp
+cfgrib

performance/fdb_performance_3D.py

@@ -20,6 +20,7 @@ def setup_method(self, method):
             "date": {"transformation": {"merge": {"with": "time", "linkers": [" ", "00"]}}},
             "step": {"transformation": {"type_change": "int"}},
             "levelist": {"transformation": {"type_change": "int"}},
+            "longitude": {"transformation": {"cyclic": [0, 360]}},
         }
         self.config = {"class": "od", "expver": "0001", "levtype": "sfc"}
         self.fdbdatacube = FDBDatacube(self.config, axis_options=self.options)

polytope/datacube/backends/datacube.py

@@ -1,12 +1,10 @@
-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 ..transformations.datacube_transformations import (
@@ -16,6 +14,24 @@
 
 
 class Datacube(ABC):
+    def __init__(self, axis_options=None, datacube_options=None):
+        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 = []
+        self.blocked_axes = []
+        self.fake_axes = []
+        self.treated_axes = []
+        self.nearest_search = {}
+        self._axes = None
+        self.transformed_axes = []
+
     @abstractmethod
     def get(self, requests: IndexTree) -> Any:
         """Return data given a set of request trees"""
@@ -46,33 +62,29 @@ def _create_axes(self, name, values, transformation_type_key, transformation_opt
 
             # 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
             # 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():
+            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):
         if options is not None:
             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 +108,12 @@ def get_indices(self, path: DatacubePath, axis, lower, upper, method=None):
         """
         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

polytope/datacube/backends/fdb.py

@@ -11,25 +11,13 @@ class FDBDatacube(Datacube):
     def __init__(self, config=None, axis_options=None, datacube_options=None, point_cloud_options=None):
         if config is None:
             config = {}
-        if axis_options is None:
-            axis_options = {}
-        if datacube_options is None:
-            datacube_options = {}
+
+        super().__init__(axis_options, datacube_options)
 
         logging.info("Created an FDB datacube with options: " + str(axis_options))
 
-        self.axis_options = axis_options
-        self.axis_counter = 0
-        self._axes = None
-        treated_axes = []
-        self.complete_axes = []
-        self.blocked_axes = []
-        self.fake_axes = []
         self.unwanted_path = {}
         self.has_point_cloud = point_cloud_options  # NOTE: here, will be True/False
-        self.coupled_axes = []
-        self.axis_with_identical_structure_after = datacube_options.get("identical structure after")
-        self.nearest_search = {}
 
         partial_request = config
         # Find values in the level 3 FDB datacube
@@ -44,12 +32,12 @@ def __init__(self, config=None, axis_options=None, datacube_options=None, point_
             values.sort()
             options = axis_options.get(name, None)
             self._check_and_add_axes(options, name, values)
-            treated_axes.append(name)
+            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 treated_axes:
+            if name not in self.treated_axes:
                 options = axis_options.get(name, None)
                 val = self._axes[name].type
                 self._check_and_add_axes(options, name, val)
@@ -250,7 +238,7 @@ def sort_fdb_request_ranges(self, range_lengths, current_start_idx, lat_length):
         return (original_indices, sorted_request_ranges)
 
     def datacube_natural_indexes(self, axis, subarray):
-        indexes = subarray[axis.name]
+        indexes = subarray.get(axis.name, None)
         return indexes
 
     def select(self, path, unmapped_path):

polytope/datacube/backends/mock.py

@@ -8,6 +8,7 @@
 
 class MockDatacube(Datacube):
     def __init__(self, dimensions, datacube_options={}):
+        super().__init__({}, datacube_options)
         assert isinstance(dimensions, dict)
 
         self.dimensions = dimensions
@@ -22,8 +23,6 @@ def __init__(self, dimensions, datacube_options={}):
         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):
         # Takes in a datacube and verifies the leaves of the tree are complete

polytope/datacube/backends/xarray.py

@@ -1,5 +1,6 @@
 from copy import deepcopy
 
+import numpy as np
 import xarray as xr
 
 from .datacube import Datacube, IndexTree
@@ -9,44 +10,32 @@ class XArrayDatacube(Datacube):
     """Xarray arrays are labelled, axes can be defined as strings or integers (e.g. "time" or 0)."""
 
     def __init__(self, dataarray: xr.DataArray, axis_options=None, datacube_options=None, point_cloud_options=None):
-        if axis_options is None:
-            axis_options = {}
-        if datacube_options is None:
-            datacube_options = {}
-        self.axis_options = axis_options
-        self.axis_counter = 0
-        self._axes = None
+        super().__init__(axis_options, datacube_options)
         self.dataarray = dataarray
         treated_axes = []
-        self.complete_axes = []
-        self.blocked_axes = []
-        self.fake_axes = []
-        self.nearest_search = None
-        self.coupled_axes = []
-        self.axis_with_identical_structure_after = datacube_options.get("identical structure after")
         self.has_point_cloud = point_cloud_options
 
         for name, values in dataarray.coords.variables.items():
             if name in dataarray.dims:
-                options = axis_options.get(name, None)
+                options = self.axis_options.get(name, None)
                 self._check_and_add_axes(options, name, values)
-                treated_axes.append(name)
+                self.treated_axes.append(name)
                 self.complete_axes.append(name)
             else:
                 if self.dataarray[name].dims == ():
-                    options = axis_options.get(name, None)
+                    options = self.axis_options.get(name, None)
                     self._check_and_add_axes(options, name, values)
-                    treated_axes.append(name)
+                    self.treated_axes.append(name)
         for name in dataarray.dims:
-            if name not in treated_axes:
-                options = axis_options.get(name, None)
+            if name not in self.treated_axes:
+                options = self.axis_options.get(name, None)
                 val = dataarray[name].values[0]
                 self._check_and_add_axes(options, name, val)
-                treated_axes.append(name)
+                self.treated_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 treated_axes:
-                options = axis_options.get(name, None)
+            if name not in self.treated_axes:
+                options = self.axis_options.get(name, None)
                 val = self._axes[name].type
                 self._check_and_add_axes(options, name, val)
 
@@ -55,45 +44,23 @@ def find_point_cloud(self):
         if self.has_point_cloud:
             return self.has_point_cloud
 
-    def old_get(self, requests: IndexTree):
-        for r in requests.leaves:
-            path = r.flatten()
-            if len(path.items()) == self.axis_counter:
-                # first, find the grid mapper transform
-                unmapped_path = {}
-                path_copy = deepcopy(path)
-                for key in path_copy:
-                    axis = self._axes[key]
-                    (path, unmapped_path) = axis.unmap_to_datacube(path, unmapped_path)
-                path = self.fit_path(path)
-                subxarray = self.dataarray.sel(path, method="nearest")
-                subxarray = subxarray.sel(unmapped_path)
-                value = subxarray.item()
-                key = subxarray.name
-                r.result = (key, value)
-            else:
-                r.remove_branch()
-
     def get(self, requests: IndexTree):
-        # TODO: change to work with the irregular grid
-        axis_counter = self.axis_counter + 1
-        # if self.has_point_cloud:
-        #     axis_counter = self.axis_counter - 1
-        # else:
-        #     axis_counter = self.axis_counter
-
         for r in requests.leaves:
-            # path = r.flatten()
             path = r.flatten_with_result()
-            if len(path.items()) == axis_counter:
+            if len(path.items()) == self.axis_counter + 1:
                 # first, find the grid mapper transform
                 unmapped_path = {}
                 path_copy = deepcopy(path)
                 for key in path_copy:
                     if key != "result":
                         axis = self._axes[key]
-                        (path, unmapped_path) = axis.unmap_to_datacube(path, unmapped_path)
-                path = self.fit_path(path)
+                        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)
+                path.update(unmapped_path)
+
+                unmapped_path = {}
+                self.refit_path(path, unmapped_path, path)
                 subxarray = self.dataarray.sel(path, method="nearest")
                 subxarray = subxarray.sel(unmapped_path)
                 value = subxarray.item()
@@ -107,13 +74,34 @@ def datacube_natural_indexes(self, axis, subarray):
             indexes = next(iter(subarray.xindexes.values())).to_pandas_index()
         else:
             if subarray[axis.name].values.ndim == 0:
-                indexes = [subarray[axis.name].values]
+                # NOTE how we handle the two special datetime and timedelta cases to conform with numpy arrays
+                if np.issubdtype(subarray[axis.name].values.dtype, np.datetime64):
+                    indexes = [subarray[axis.name].astype("datetime64[us]").values]
+                elif np.issubdtype(subarray[axis.name].values.dtype, np.timedelta64):
+                    indexes = [subarray[axis.name].astype("timedelta64[us]").values]
+                else:
+                    indexes = [subarray[axis.name].values.tolist()]
             else:
                 indexes = subarray[axis.name].values
         return indexes
 
+    def refit_path(self, path_copy, unmapped_path, path):
+        for key in path.keys():
+            if key not in self.dataarray.dims:
+                path_copy.pop(key)
+            elif key not in self.dataarray.coords.dtypes:
+                unmapped_path.update({key: path[key]})
+                path_copy.pop(key, None)
+            for key in self.dataarray.coords.dtypes:
+                key_dtype = self.dataarray.coords.dtypes[key]
+                if key_dtype.type is np.str_ and key in path.keys():
+                    unmapped_path.update({key: path[key]})
+                    path_copy.pop(key, None)
+
     def select(self, path, unmapped_path):
-        subarray = self.dataarray.sel(path, method="nearest")
+        path_copy = deepcopy(path)
+        self.refit_path(path_copy, unmapped_path, path)
+        subarray = self.dataarray.sel(path_copy, method="nearest")
         subarray = subarray.sel(unmapped_path)
         return subarray