Work on indexing

Author: b8raoult

ecml_tools/data.py

@@ -165,6 +165,10 @@ def metadata_specific(self, **kwargs):
     def __repr__(self):
         return self.__class__.__name__ + "()"
 
+    def _get_tuple(self, n):
+        first, rest = n[0], n[1:]
+        return self[first][rest]
+
 
 class Source:
     def __init__(self, dataset, index, source=None, info=None):
@@ -287,13 +291,17 @@ def __len__(self):
     def __getitem__(self, n):
         return self.data[n]
 
+    @cached_property
+    def chunks(self):
+        return self.z.data.chunks
+
     @cached_property
     def shape(self):
         return self.data.shape
 
     @cached_property
     def dtype(self):
-        return self.data.dtype
+        return self.z.data.dtype
 
     @cached_property
     def dates(self):
@@ -361,7 +369,11 @@ def end_of_statistics_date(self):
         return self.dates[-1]
 
     def metadata_specific(self):
-        return super().metadata_specific(attrs=dict(self.z.attrs))
+        return super().metadata_specific(
+            attrs=dict(self.z.attrs),
+            chunks=self.chunks,
+            dtype=str(self.dtype),
+        )
 
     def source(self, index):
         return Source(self, index, info=self.path)
@@ -528,6 +540,9 @@ def __len__(self):
         return sum(len(i) for i in self.datasets)
 
     def __getitem__(self, n):
+        if isinstance(n, tuple):
+            return self._get_tuple(n)
+
         if isinstance(n, slice):
             return self._get_slice(n)
 
@@ -609,8 +624,12 @@ def _get_slice(self, s):
         return np.stack([self[i] for i in range(*s.indices(self._len))])
 
     def __getitem__(self, n):
+        if isinstance(n, tuple):
+            return self._get_tuple(n)
+
         if isinstance(n, slice):
             return self._get_slice(n)
+
         return np.concatenate([d[n] for d in self.datasets], axis=self.axis - 1)
 
 
@@ -658,8 +677,12 @@ def _get_slice(self, s):
         return np.stack([self[i] for i in range(*s.indices(self._len))])
 
     def __getitem__(self, n):
+        if isinstance(n, tuple):
+            return self._get_tuple(n)
+
         if isinstance(n, slice):
             return self._get_slice(n)
+
         return np.concatenate([d[n] for d in self.datasets])
 
     @cached_property
@@ -743,8 +766,12 @@ def __init__(self, dataset, indices):
         super().__init__(dataset)
 
     def __getitem__(self, n):
+        if isinstance(n, tuple):
+            return self._get_tuple(n)
+
         if isinstance(n, slice):
             return self._get_slice(n)
+
         n = self.indices[n]
         return self.dataset[n]
 
@@ -794,9 +821,13 @@ def __init__(self, dataset, indices):
         super().__init__(dataset)
 
     def __getitem__(self, n):
+        if isinstance(n, tuple):
+            return self._get_tuple(n)
+
         row = self.dataset[n]
         if isinstance(n, slice):
             return row[:, self.indices]
+
         return row[self.indices]
 
     @cached_property

tests/test_data.py

@@ -66,28 +66,54 @@ def create_zarr(
     ensembles = ensemble if ensemble is not None else 1
     values = grids if grids is not None else VALUES
 
-    data = np.zeros((len(dates), len(vars), ensembles, values))
+    data = np.zeros(shape=(len(dates), len(vars), ensembles, values))
 
     for i, date in enumerate(dates):
         for j, var in enumerate(vars):
             for e in range(ensembles):
                 data[i, j, e] = _(date.astype(object), var, k, e, values)
 
-    root.data = data
-    root.dates = dates
-    root.latitudes = np.array([x + values for x in range(values)])
-    root.longitudes = np.array([x + values for x in range(values)])
+    root.create_dataset(
+        "data",
+        data=data,
+        dtype=data.dtype,
+        chunks=data.shape,
+    )
+    root.create_dataset(
+        "dates",
+        data=dates,
+    )
+    root.create_dataset(
+        "latitudes",
+        data=np.array([x + values for x in range(values)]),
+    )
+    root.create_dataset(
+        "longitudes",
+        data=np.array([x + values for x in range(values)]),
+    )
 
     root.attrs["frequency"] = frequency
     root.attrs["resolution"] = resolution
     root.attrs["name_to_index"] = {k: i for i, k in enumerate(vars)}
 
     root.attrs["data_request"] = {"grid": 1, "area": "g", "param_level": {}}
 
-    root.mean = np.mean(data, axis=0)
-    root.stdev = np.std(data, axis=0)
-    root.maximum = np.max(data, axis=0)
-    root.minimum = np.min(data, axis=0)
+    root.create_dataset(
+        "mean",
+        data=np.mean(data, axis=0),
+    )
+    root.create_dataset(
+        "stdev",
+        data=np.std(data, axis=0),
+    )
+    root.create_dataset(
+        "maximum",
+        data=np.max(data, axis=0),
+    )
+    root.create_dataset(
+        "minimum",
+        data=np.min(data, axis=0),
+    )
 
     return root
 
@@ -167,6 +193,11 @@ def slices(ds, start=None, end=None, step=None):
     for i, n in enumerate(range(start, end, step)):
         assert (s[i] == ds[n]).all()
 
+    x = ds[0:10, :, 0]
+
+    if ds.shape[2] > 1:
+        ds[0:10, :, np.array([1, 0])]
+
 
 def make_row(args, ensemble=False, grid=False):
     if grid: