From a244b524dfcb4898899304e176f1c932f0fe079c Mon Sep 17 00:00:00 2001 From: mathleur Date: Mon, 4 Nov 2024 20:27:04 +0100 Subject: [PATCH 1/2] pass context to gj axes and move context --- polytope_feature/datacube/backends/datacube.py | 6 +++--- polytope_feature/datacube/backends/fdb.py | 6 ++++-- polytope_feature/datacube/backends/xarray.py | 2 +- polytope_feature/polytope.py | 18 +++++++++--------- 4 files changed, 17 insertions(+), 15 deletions(-) diff --git a/polytope_feature/datacube/backends/datacube.py b/polytope_feature/datacube/backends/datacube.py index 58c186a6..7f43c6b6 100644 --- a/polytope_feature/datacube/backends/datacube.py +++ b/polytope_feature/datacube/backends/datacube.py @@ -148,17 +148,17 @@ def remap_path(self, path: DatacubePath): return path @staticmethod - def create(datacube, config={}, axis_options={}, compressed_axes_options=[], alternative_axes=[]): + def create(datacube, config={}, axis_options={}, compressed_axes_options=[], alternative_axes=[], context=None): # 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, compressed_axes_options) + xadatacube = XArrayDatacube(datacube, axis_options, compressed_axes_options, context) return xadatacube if type(datacube).__name__ == "GribJump": from .fdb import FDBDatacube - fdbdatacube = FDBDatacube(datacube, config, axis_options, compressed_axes_options, alternative_axes) + fdbdatacube = FDBDatacube(datacube, config, axis_options, compressed_axes_options, alternative_axes, context) return fdbdatacube def check_branching_axes(self, request): diff --git a/polytope_feature/datacube/backends/fdb.py b/polytope_feature/datacube/backends/fdb.py index cc427f56..714e741c 100644 --- a/polytope_feature/datacube/backends/fdb.py +++ b/polytope_feature/datacube/backends/fdb.py @@ -8,9 +8,11 @@ class FDBDatacube(Datacube): - def __init__(self, gj, config=None, axis_options=None, compressed_axes_options=[], alternative_axes=[]): + def __init__(self, gj, config=None, axis_options=None, compressed_axes_options=[], alternative_axes=[], context=None): if config is None: config = {} + if context is None: + context = {} super().__init__(axis_options, compressed_axes_options) @@ -24,7 +26,7 @@ def __init__(self, gj, config=None, axis_options=None, compressed_axes_options=[ self.gj = gj if len(alternative_axes) == 0: - self.fdb_coordinates = self.gj.axes(partial_request) + self.fdb_coordinates = self.gj.axes(partial_request, ctx=context) else: self.fdb_coordinates = {} for axis_config in alternative_axes: diff --git a/polytope_feature/datacube/backends/xarray.py b/polytope_feature/datacube/backends/xarray.py index 2b7d579d..735c87e2 100644 --- a/polytope_feature/datacube/backends/xarray.py +++ b/polytope_feature/datacube/backends/xarray.py @@ -9,7 +9,7 @@ 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, compressed_axes_options=[]): + def __init__(self, dataarray: xr.DataArray, axis_options=None, compressed_axes_options=[], context=None): super().__init__(axis_options, compressed_axes_options) if axis_options is None: axis_options = {} diff --git a/polytope_feature/polytope.py b/polytope_feature/polytope.py index 29a03d63..2d0a1f35 100644 --- a/polytope_feature/polytope.py +++ b/polytope_feature/polytope.py @@ -39,7 +39,7 @@ def __repr__(self): class Polytope: - def __init__(self, datacube, engine=None, options=None): + def __init__(self, datacube, engine=None, options=None, context=None): from .datacube import Datacube from .engine import Engine @@ -48,7 +48,9 @@ def __init__(self, datacube, engine=None, options=None): axis_options, compressed_axes_options, config, alternative_axes = PolytopeOptions.get_polytope_options(options) - self.datacube = Datacube.create(datacube, config, axis_options, compressed_axes_options, alternative_axes) + self.context = context + + self.datacube = Datacube.create(datacube, config, axis_options, compressed_axes_options, alternative_axes, self.context) self.engine = engine if engine is not None else Engine.default() self.time = 0 @@ -56,14 +58,12 @@ def slice(self, polytopes: List[ConvexPolytope]): """Low-level API which takes a polytope geometry object and uses it to slice the datacube""" return self.engine.extract(self.datacube, polytopes) - def retrieve(self, request: Request, method="standard", context=None): + def retrieve(self, request: Request, method="standard"): """Higher-level API which takes a request and uses it to slice the datacube""" - if context is None: - context = {} - logging.info("Starting request for %s ", context) + logging.info("Starting request for %s ", self.context) self.datacube.check_branching_axes(request) request_tree = self.engine.extract(self.datacube, request.polytopes()) - logging.info("Created request tree for %s ", context) - self.datacube.get(request_tree, context) - logging.info("Retrieved data for %s ", context) + logging.info("Created request tree for %s ", self.context) + self.datacube.get(request_tree, self.context) + logging.info("Retrieved data for %s ", self.context) return request_tree From 684abc1c7c3461bf310fa1be4152ca5924cc3b3a Mon Sep 17 00:00:00 2001 From: mathleur Date: Mon, 4 Nov 2024 20:28:10 +0100 Subject: [PATCH 2/2] black --- polytope_feature/datacube/backends/datacube.py | 4 +++- polytope_feature/datacube/backends/fdb.py | 4 +++- polytope_feature/polytope.py | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/polytope_feature/datacube/backends/datacube.py b/polytope_feature/datacube/backends/datacube.py index 7f43c6b6..1faebf19 100644 --- a/polytope_feature/datacube/backends/datacube.py +++ b/polytope_feature/datacube/backends/datacube.py @@ -158,7 +158,9 @@ def create(datacube, config={}, axis_options={}, compressed_axes_options=[], alt if type(datacube).__name__ == "GribJump": from .fdb import FDBDatacube - fdbdatacube = FDBDatacube(datacube, config, axis_options, compressed_axes_options, alternative_axes, context) + fdbdatacube = FDBDatacube( + datacube, config, axis_options, compressed_axes_options, alternative_axes, context + ) return fdbdatacube def check_branching_axes(self, request): diff --git a/polytope_feature/datacube/backends/fdb.py b/polytope_feature/datacube/backends/fdb.py index 714e741c..138f8b6a 100644 --- a/polytope_feature/datacube/backends/fdb.py +++ b/polytope_feature/datacube/backends/fdb.py @@ -8,7 +8,9 @@ class FDBDatacube(Datacube): - def __init__(self, gj, config=None, axis_options=None, compressed_axes_options=[], alternative_axes=[], context=None): + def __init__( + self, gj, config=None, axis_options=None, compressed_axes_options=[], alternative_axes=[], context=None + ): if config is None: config = {} if context is None: diff --git a/polytope_feature/polytope.py b/polytope_feature/polytope.py index 2d0a1f35..f18f10c4 100644 --- a/polytope_feature/polytope.py +++ b/polytope_feature/polytope.py @@ -50,7 +50,9 @@ def __init__(self, datacube, engine=None, options=None, context=None): self.context = context - self.datacube = Datacube.create(datacube, config, axis_options, compressed_axes_options, alternative_axes, self.context) + self.datacube = Datacube.create( + datacube, config, axis_options, compressed_axes_options, alternative_axes, self.context + ) self.engine = engine if engine is not None else Engine.default() self.time = 0