refactor complete; no tests

Author: rabernat

xbatcher/generators.py

@@ -11,6 +11,52 @@ def _as_xarray_dataset(ds):
     else:
         return ds.to_dataset()
 
+def _slices(dimsize, size, overlap=0):
+    # return a list of slices to chop up a single dimension
+    slices = []
+    stride = size - overlap
+    assert stride > 0
+    assert stride < dimsize
+    for start in range(0, dimsize, stride):
+        end = start+size
+        if end <= dimsize:
+            slices.append(slice(start, end))
+    return slices
+
+
+def _iterate_through_dataset(ds, dims, overlap={}):
+    dim_slices = []
+    for dim in dims:
+        dimsize = ds.dims[dim]
+        size = dims[dim]
+        olap = overlap.get(dim, 0)
+        dim_slices.append(_slices(dimsize, size, olap))
+
+    for slices in itertools.product(*dim_slices):
+        selector = {key: slice for key, slice in zip(dims, slices)}
+        yield ds.isel(**selector)
+
+
+def _drop_input_dims(ds, input_dims, suffix='_input'):
+    # remove input_dims coordinates from datasets, rename the dimensions
+    # then put intput_dims back in as coordinates
+    out = ds.copy()
+    out = (out.drop(input_dims)
+              .rename({dim: dim + suffix for dim in input_dims}))
+    for dim in input_dims:
+        out.coords[dim] = dim + suffix, ds[dim].values
+    return out
+
+
+def _maybe_stack_batch_dims(ds, input_dims):
+    batch_dims = list(set(ds.dims) - set(input_dims))
+    if len(batch_dims) < 2:
+        return ds
+    ds_stack = ds.stack(batch=batch_dims)
+    # ensure correct order
+    dim_order = ('batch',) + tuple(input_dims)
+    return ds_stack.transpose(*dim_order)
+
 
 class BatchGenerator:
     """Create generator for iterating through xarray datarrays / datasets in
@@ -20,44 +66,59 @@ class BatchGenerator:
     ----------
     ds : ``xarray.Dataset`` or ``xarray.DataArray``
         The data to iterate over
-    batch_sizes : dict
-        A dictionary specifying the size of the batch in each dimension,
-        e.g. ``{'time': 100, 'latitude': 30}``
-    overlap : dict, optional
+    input_dims : dict
+        A dictionary specifying the size of the inputs in each dimension,
+        e.g. ``{'lat': 30, 'lon': 30}``
+        These are the dimensions the ML library will see. All other dimensions
+        will be stacked into one dimension called ``batch``.
+    input_overlap : dict, optional
         A dictionary specifying the overlap along each dimension
+        e.g. ``{'lat': 3, 'lon': 3}``
+    batch_dims : dict, optional
+        A dictionary specifying the size of the batch along each dimension
+        e.g. ``{'time': 10}``. These will always be interated over.
+    concat_input_dims : bool, optional
+        If ``True``, the dimension chunks specified in ``input_dims`` will be
+        concatenated and stacked into the batch dimension. If ``False``, they
+        will be iterated over.
+    preload_batch : bool, optional
+        If ``True``, each batch will be loaded into memory before reshaping /
+        processing, triggering any dask arrays to be computed.
 
     Yields
     ------
     ds_slice : ``xarray.Dataset`` or ``xarray.DataArray``
-        Slices of the array matching the given batch size specification
+        Slices of the array matching the given batch size specification.
     """
 
-    def __init__(self, ds, batch_sizes, overlap={}):
+    def __init__(self, ds, input_dims, input_overlap={}, batch_dims={},
+                 concat_input_dims=False, preload_batch=True):
+
         self.ds = _as_xarray_dataset(ds)
         # should be a dict
-        self.batch_sizes = OrderedDict(batch_sizes)
-        self.batch_dims = list(self.batch_sizes)
-        # make overlap is defined for each batch size defined
-        self.overlap = {k: overlap.get(k, 0) for k in self.batch_dims}
+        self.input_dims = OrderedDict(input_dims)
+        self.input_overlap = input_overlap
+        self.batch_dims = OrderedDict(batch_dims)
+        self.concat_input_dims = concat_input_dims
+        self.preload_batch = preload_batch
 
 
     def __iter__(self):
-        for slices in itertools.product(*[self._iterate_dim(dim)
-                                          for dim in self.batch_dims]):
-            selector = {key: slice for key, slice in zip(self.batch_dims, slices)}
-            yield self.ds.isel(**selector)
-
-
-    def _iterate_dim(self, dim):
-        dimsize = self.ds.dims[dim]
-        size = self.batch_sizes[dim]
-        overlap = self.overlap[dim]
-        stride = size - overlap
-        assert stride > 0
-        assert stride < dimsize
-        for start in range(0, dimsize, stride):
-            end = start+size
-            if end <= dimsize:
-                yield slice(start, end)
+        for ds_batch in self._iterate_batch_dims(self.ds):
+            if self.preload_batch:
+                ds_batch.load()
+            input_generator = self._iterate_input_dims(ds_batch)
+            if self.concat_input_dims:
+                all_dsets = [_drop_input_dims(ds_input, list(self.input_dims))
+                             for ds_input in input_generator]
+                dsc = xr.concat(all_batches, dim='input_batch')
+                yield _maybe_stack_batch_dims(dsc, list(self.input_dims))
             else:
-                return
+                for ds_input in input_generator:
+                    yield _maybe_stack_batch_dims(ds_input, list(self.input_dims))
+
+    def _iterate_batch_dims(self, ds):
+        return _iterate_through_dataset(ds, self.batch_dims)
+
+    def _iterate_input_dims(self, ds):
+        return _iterate_through_dataset(ds, self.input_dims, self.input_overlap)

xbatcher/tests/test_generators.py

@@ -17,7 +17,7 @@ def sample_ds_1d():
 # Should we enforce that each batch size always has to be the same
 @pytest.mark.parametrize("bsize", [5, 10])
 def test_batch_1d(sample_ds_1d, bsize):
-    bg = BatchGenerator(sample_ds_1d, batch_sizes={'x': bsize})
+    bg = BatchGenerator(sample_ds_1d, input_dims={'x': bsize})
     for n, ds_batch in enumerate(bg):
         assert isinstance(ds_batch, xr.Dataset)
         # TODO: maybe relax this? see comment above
@@ -30,8 +30,8 @@ def test_batch_1d(sample_ds_1d, bsize):
 @pytest.mark.parametrize("olap", [1, 4])
 def test_batch_1d_overlap(sample_ds_1d, olap):
     bsize = 10
-    bg = BatchGenerator(sample_ds_1d, batch_sizes={'x': bsize},
-                        overlap={'x': olap})
+    bg = BatchGenerator(sample_ds_1d, input_dims={'x': bsize},
+                        input_overlap={'x': olap})
     stride = bsize-olap
     for n, ds_batch in enumerate(bg):
         assert isinstance(ds_batch, xr.Dataset)
@@ -55,7 +55,7 @@ def sample_ds_2d():
 def test_batch_2d(sample_ds_2d, bsize):
 
     # first do the iteration over just one dimension
-    bg = BatchGenerator(sample_ds_2d, batch_sizes={'x': bsize})
+    bg = BatchGenerator(sample_ds_2d, input_dims={'x': bsize})
     for n, ds_batch in enumerate(bg):
         assert isinstance(ds_batch, xr.Dataset)
         assert ds_batch.dims['x'] == bsize
@@ -66,7 +66,7 @@ def test_batch_2d(sample_ds_2d, bsize):
 
     # now iterate over both
     xbsize = 20
-    bg = BatchGenerator(sample_ds_2d, batch_sizes={'y': bsize, 'x': xbsize})
+    bg = BatchGenerator(sample_ds_2d, input_dims={'y': bsize, 'x': xbsize})
     for n, ds_batch in enumerate(bg):
         assert isinstance(ds_batch, xr.Dataset)
         assert ds_batch.dims['x'] == xbsize