first commit

Author: rabernat

.gitignore

@@ -0,0 +1,63 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# notebook
+*/.ipynb_checkpoints/*
+
+# tests
+.pytest_cache/*

setup.py

@@ -0,0 +1,50 @@
+ #!/usr/bin/env python
+import os
+import re
+import sys
+import warnings
+
+from setuptools import setup, find_packages
+
+VERSION = '0.1.0'
+DISTNAME = 'xbatcher'
+LICENSE = 'Apache'
+AUTHOR = 'xbatcher Developers'
+AUTHOR_EMAIL = '[email protected]'
+URL = 'https://github.com/xgcm/xbatcher'
+CLASSIFIERS = [
+    'Development Status :: 4 - Beta',
+    'License :: OSI Approved :: Apache Software License',
+    'Operating System :: OS Independent',
+    'Intended Audience :: Science/Research',
+    'Programming Language :: Python',
+    'Programming Language :: Python :: 3',
+    'Programming Language :: Python :: 3.5',
+    'Programming Language :: Python :: 3.6',
+    'Topic :: Scientific/Engineering',
+]
+
+INSTALL_REQUIRES = ['xarray', 'dask', 'numpy']
+SETUP_REQUIRES = []
+TESTS_REQUIRE = ['pytest >= 2.8', 'coverage']
+
+DESCRIPTION = "Batch generation from xarray dataset"
+def readme():
+    return "TODO"
+    #with open('README.rst') as f:
+    #    return f.read()
+
+
+setup(name=DISTNAME,
+      version=VERSION,
+      license=LICENSE,
+      author=AUTHOR,
+      author_email=AUTHOR_EMAIL,
+      classifiers=CLASSIFIERS,
+      description=DESCRIPTION,
+      long_description=readme(),
+      install_requires=INSTALL_REQUIRES,
+      setup_requires=SETUP_REQUIRES,
+      tests_require=TESTS_REQUIRE,
+      url=URL,
+      packages=find_packages())

xbatcher/__init__.py

@@ -0,0 +1 @@
+from . generators import BatchGenerator

xbatcher/generators.py

@@ -0,0 +1,44 @@
+import xarray as xr
+from collections import OrderedDict
+import itertools
+
+def _as_xarray_dataset(ds):
+    # maybe coerce to xarray dataset
+    if isinstance(ds, xr.Dataset):
+        return ds
+    else:
+        return ds.to_dataset()
+
+
+class BatchGenerator:
+    """For iterating through xarray datarrays / datasets in batches."""
+
+    def __init__(self, ds, batch_sizes, overlap={}):
+        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}
+
+
+    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)
+            else:
+                return

xbatcher/tests/test_generators.py

@@ -0,0 +1,39 @@
+import xarray as xr
+import numpy as np
+from xbatcher import BatchGenerator
+import pytest
+
+
+@pytest.fixture(scope='module')
+def sample_ds_1d():
+    size=100
+    ds = xr.Dataset({'foo': (['x'], np.random.rand(size)),
+                     'bar': (['x'], np.random.randint(0, 10, size))},
+                    {'x': (['x'], np.arange(size))})
+    return ds
+
+# TODO: decide how to handle bsizes like 15 that don't evenly divide the dimension
+# 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})
+    for n, ds_batch in enumerate(bg):
+        assert isinstance(ds_batch, xr.Dataset)
+        # TODO: maybe relax this? see comment above
+        assert ds_batch.dims['x'] == bsize
+        expected_slice = slice(bsize*n, bsize*(n+1))
+        ds_batch_expected = sample_ds_1d.isel(x=expected_slice)
+        assert ds_batch.equals(ds_batch_expected)
+
+@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})
+    stride = bsize-olap
+    for n, ds_batch in enumerate(bg):
+        assert isinstance(ds_batch, xr.Dataset)
+        assert ds_batch.dims['x'] == bsize
+        expected_slice = slice(stride*n, stride*n + bsize)
+        ds_batch_expected = sample_ds_1d.isel(x=expected_slice)
+        assert ds_batch.equals(ds_batch_expected)