Merge pull request #111 from ecmwf/feature/reformat_datacube_init

add a global datacube init

polytope/datacube/backends/datacube.py

@@ -16,6 +16,25 @@
 
 
 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"""

polytope/datacube/backends/fdb.py

@@ -11,26 +11,12 @@ class FDBDatacube(Datacube):
     def __init__(self, config=None, axis_options=None, datacube_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.nearest_search = {}
-        self.coupled_axes = []
-        self.axis_with_identical_structure_after = datacube_options.get("identical structure after")
-        self.transformed_axes = []
-        self.dataarray = []
 
         partial_request = config
         # Find values in the level 3 FDB datacube
@@ -45,12 +31,12 @@ def __init__(self, config=None, axis_options=None, datacube_options=None):
             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)

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

@@ -10,44 +10,30 @@ 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):
-        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.transformed_axes = []
 
         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)