Indexing test

Author: b8raoult

ecml_tools/data.py

@@ -191,9 +191,7 @@ def __repr__(self):
         return self.__class__.__name__ + "()"
 
     def _get_tuple(self, n):
-        warnings.warn(f"Naive tuple indexing used with {self}, likely to be slow.")
-        first, rest = n[0], n[1:]
-        return self[first][rest]
+        raise NotImplementedError(f"Tuple not supported: {n} (class {self.__class__.__name__})")
 
 
 class Source:
@@ -932,11 +930,11 @@ def __init__(self, dataset, indices):
         super().__init__(dataset)
 
     def __getitem__(self, n):
-        if isinstance(n, tuple):
-            return self._get_tuple(n)
+        # if isinstance(n, tuple):
+        #     return self._get_tuple(n)
 
         row = self.dataset[n]
-        if isinstance(n, slice):
+        if isinstance(n, (slice, tuple)):
             return row[:, self.indices]
 
         return row[self.indices]

tests/test_data.py

@@ -173,6 +173,15 @@ def same_stats(ds1, ds2, vars1, vars2=None):
         ).all()
 
 
+class IndexTester:
+    def __init__(self, ds):
+        self.ds = ds
+        self.np = ds[:]  # Numpy array
+
+    def __getitem__(self, index):
+        assert (self.ds[index] == self.np[index]).all()
+
+
 def slices(ds, start=None, end=None, step=None):
     if start is None:
         start = 5
@@ -181,22 +190,18 @@ def slices(ds, start=None, end=None, step=None):
     if step is None:
         step = len(ds) // 10
 
-    s = ds[start:end:step]
-
-    assert s[0].shape == ds[0].shape, (
-        s.shape,
-        ds.shape,
-        len(list(range(start, end, step))),
-        list(range(start, end, step)),
-    )
+    t = IndexTester(ds)
 
-    for i, n in enumerate(range(start, end, step)):
-        assert (s[i] == ds[n]).all()
+    t[start:end:step]
+    t[start:end]
+    t[start:]
+    t[:end]
+    t[::step]
 
-    ds[0:10, :, 0]
+    t[0:10, :, 0]
 
-    if ds.shape[2] > 1:
-        ds[0:10, :, np.array([1, 0])]
+    if ds.shape[2] > 1:  # Ensemble dimension
+        t[0:10, :, (0, 1)]
 
 
 def make_row(args, ensemble=False, grid=False):