update

Author: b8raoult

ecml_tools/data.py

@@ -33,6 +33,7 @@
 __all__ = ["open_dataset", "open_zarr", "debug_zarr_loading"]
 
 DEBUG_ZARR_LOADING = int(os.environ.get("DEBUG_ZARR_LOADING", "0"))
+DEBUG_ZARR_INDEXING = int(os.environ.get("DEBUG_ZARR_INDEXING", "0"))
 
 DEPTH = 0
 
@@ -41,25 +42,22 @@ def _debug_indexing(method):
     @wraps(method)
     def wrapper(self, index):
         global DEPTH
-        if isinstance(index, tuple):
-            print("  " * DEPTH, "->", self, method.__name__, index)
+        # if isinstance(index, tuple):
+        print("  " * DEPTH, "->", self, method.__name__, index)
         DEPTH += 1
         result = method(self, index)
         DEPTH -= 1
-        if isinstance(index, tuple):
-            print("  " * DEPTH, "<-", self, method.__name__, result.shape)
+        # if isinstance(index, tuple):
+        print("  " * DEPTH, "<-", self, method.__name__, result.shape)
         return result
 
     return wrapper
 
 
-if True:
-
-    def debug_indexing(x):
-        return x
-
-else:
+if DEBUG_ZARR_INDEXING:
     debug_indexing = _debug_indexing
+else:
+    debug_indexing = lambda x: x  # noqa
 
 
 def debug_zarr_loading(on_off):
@@ -917,9 +915,6 @@ def __init__(self, dataset, indices):
 
         self.dataset = dataset
         self.indices = list(indices)
-        self.slice = _make_slice_or_index_from_list_or_tuple(self.indices)
-        assert isinstance(self.slice, slice)
-        print("SUBSET", self.slice)
 
         # Forward other properties to the super dataset
         super().__init__(dataset)
@@ -947,7 +942,10 @@ def _get_slice(self, s):
     @expand_list_indexing
     def _get_tuple(self, n):
         index, changes = index_to_slices(n, self.shape)
-        index, previous = update_tuple(index, 0, self.slice)
+        indices = [self.indices[i] for i in range(*index[0].indices(self._len))]
+        index, previous = update_tuple(
+            index, 0, _make_slice_or_index_from_list_or_tuple(indices)
+        )
         result = self.dataset[index]
         result = result[previous]
         result = apply_index_to_slices_changes(result, changes)

tests/test_data.py

@@ -1041,6 +1041,42 @@ def test_ensemble_2():
     )
 
 
+def test_ensemble_3():
+    test = DatasetTester(
+        ensemble=[
+            {"dataset": "test-2021-2021-6h-o96-abcd-1-10", "frequency": 12},
+            {"dataset": "test-2021-2021-6h-o96-abcd-2-1", "frequency": 12},
+            {"dataset": "test-2021-2021-6h-o96-abcd-3-5", "frequency": 12},
+        ]
+    )
+    test.run(
+        expected_class=Ensemble,
+        expected_length=365 * 1 * 2,
+        expected_shape=(365 * 1 * 2, 4, 16, VALUES),
+        expected_variables="abcd",
+        expected_name_to_index={"a": 0, "b": 1, "c": 2, "d": 3},
+        date_to_row=lambda date: make_row(
+            [_(date, "a", 1, i) for i in range(10)]
+            + [_(date, "a", 2, 0)]
+            + [_(date, "a", 3, i) for i in range(5)],
+            [_(date, "b", 1, i) for i in range(10)]
+            + [_(date, "b", 2, 0)]
+            + [_(date, "b", 3, i) for i in range(5)],
+            [_(date, "c", 1, i) for i in range(10)]
+            + [_(date, "c", 2, 0)]
+            + [_(date, "c", 3, i) for i in range(5)],
+            [_(date, "d", 1, i) for i in range(10)]
+            + [_(date, "d", 2, 0)]
+            + [_(date, "d", 3, i) for i in range(5)],
+            ensemble=True,
+        ),
+        start_date=datetime.datetime(2021, 1, 1),
+        time_increment=datetime.timedelta(hours=12),
+        statistics_reference_dataset=None,
+        statistics_reference_variables=None,
+    )
+
+
 def test_grids():
     test = DatasetTester(
         grids=[