Repository: rasterio/rasterio

Issue:
Add block-wise calculation option to rio-calc
So that very large datasets can be read and written.


Discussion:

Resolving patch:
diff --git a/rasterio/rio/calc.py b/rasterio/rio/calc.py
--- a/rasterio/rio/calc.py
+++ b/rasterio/rio/calc.py
@@ -1,7 +1,10 @@
 """$ rio calc"""
 
+from __future__ import division
+
 from collections import OrderedDict
 from distutils.version import LooseVersion
+import math
 
 import click
 import snuggs
@@ -9,19 +12,20 @@
 import rasterio
 from rasterio.features import sieve
 from rasterio.fill import fillnodata
+from rasterio.windows import Window
 from rasterio.rio import options
 from rasterio.rio.helpers import resolve_inout
 
 
-def get_bands(inputs, d, i=None):
+def _get_bands(inputs, sources, d, i=None):
     """Get a rasterio.Band object from calc's inputs"""
-    path = inputs[d] if d in dict(inputs) else inputs[int(d) - 1][1]
-    src = rasterio.open(path)
+    idx = d if d in dict(inputs) else int(d) - 1
+    src = sources[idx]
     return (rasterio.band(src, i) if i else
             [rasterio.band(src, j) for j in src.indexes])
 
 
-def read_array(ix, subix=None, dtype=None):
+def _read_array(ix, subix=None, dtype=None):
     """Change the type of a read array"""
     arr = snuggs._ctx.lookup(ix, subix)
     if dtype:
@@ -29,6 +33,49 @@ def read_array(ix, subix=None, dtype=None):
     return arr
 
 
+def _chunk_output(width, height, count, itemsize, mem_limit=1):
+    """Divide the calculation output into chunks
+
+    This function determines the chunk size such that an array of shape
+    (chunk_size, chunk_size, count) with itemsize bytes per element
+    requires no more than mem_limit megabytes of memory.
+
+    Output chunks are described by rasterio Windows.
+
+    Parameters
+    ----------
+    width : int
+        Output width
+    height : int
+        Output height
+    count : int
+        Number of output bands
+    itemsize : int
+        Number of bytes per pixel
+    mem_limit : int, default
+        The maximum size in memory of a chunk array
+
+    Returns
+    -------
+    sequence of Windows
+    """
+    max_pixels = mem_limit * 1.0e+6 / itemsize * count
+    chunk_size = int(math.floor(math.sqrt(max_pixels)))
+    ncols = int(math.ceil(width / chunk_size))
+    nrows = int(math.ceil(height / chunk_size))
+    chunk_windows = []
+
+    for col in range(ncols):
+        col_offset = col * chunk_size
+        w = min(chunk_size, width - col_offset)
+        for row in range(nrows):
+            row_offset = row * chunk_size
+            h = min(chunk_size, height - row_offset)
+            chunk_windows.append(((row, col), Window(col_offset, row_offset, w, h)))
+
+    return chunk_windows
+
+
 @click.command(short_help="Raster data calculator.")
 @click.argument('command')
 @options.files_inout_arg
@@ -40,10 +87,10 @@ def read_array(ix, subix=None, dtype=None):
 @options.dtype_opt
 @options.masked_opt
 @options.overwrite_opt
+@click.option("--mem-limit", type=int, default=64, help="Limit on memory used to perform calculations, in MB.")
 @options.creation_options
 @click.pass_context
-def calc(ctx, command, files, output, name, dtype, masked, overwrite,
-         creation_options):
+def calc(ctx, command, files, output, name, dtype, masked, overwrite, mem_limit, creation_options):
     """A raster data calculator
 
     Evaluates an expression using input datasets and writes the result
@@ -56,13 +103,13 @@ def calc(ctx, command, files, output, name, dtype, masked, overwrite,
 
     \b
         * (read i) evaluates to the i-th input dataset (a 3-D array).
-        * (read i j) evaluates to the j-th band of the i-th dataset (a 2-D
-          array).
-        * (take foo j) evaluates to the j-th band of a dataset named foo (see
-          help on the --name option above).
+        * (read i j) evaluates to the j-th band of the i-th dataset (a
+          2-D array).
+        * (take foo j) evaluates to the j-th band of a dataset named foo
+          (see help on the --name option above).
         * Standard numpy array operators (+, -, *, /) are available.
-        * When the final result is a list of arrays, a multi band output
-          file is written.
+        * When the final result is a list of arrays, a multiple band
+          output file is written.
         * When the final result is a single array, a single band output
           file is written.
 
@@ -72,15 +119,18 @@ def calc(ctx, command, files, output, name, dtype, masked, overwrite,
          $ rio calc "(+ 2 (* 0.95 (read 1)))" tests/data/RGB.byte.tif \\
          > /tmp/out.tif
 
-    Produces a 3-band GeoTIFF with all values scaled by 0.95 and
-    incremented by 2.
+    The command above produces a 3-band GeoTIFF with all values scaled
+    by 0.95 and incremented by 2.
 
     \b
         $ rio calc "(asarray (+ 125 (read 1)) (read 1) (read 1))" \\
         > tests/data/shade.tif /tmp/out.tif
 
-    Produces a 3-band RGB GeoTIFF, with red levels incremented by 125,
-    from the single-band input.
+    The command above produces a 3-band RGB GeoTIFF, with red levels
+    incremented by 125, from the single-band input.
+
+    The maximum amount of memory used to perform caculations defaults to
+    64 MB. This number can be increased to improve speed of calculation.
 
     """
     import numpy as np
@@ -89,19 +139,36 @@ def calc(ctx, command, files, output, name, dtype, masked, overwrite,
         with ctx.obj['env']:
             output, files = resolve_inout(files=files, output=output,
                                           overwrite=overwrite)
-
             inputs = ([tuple(n.split('=')) for n in name] +
                       [(None, n) for n in files])
+            sources = [rasterio.open(path) for name, path in inputs]
 
-            with rasterio.open(inputs[0][1]) as first:
-                kwargs = first.meta
-                kwargs.update(**creation_options)
-                dtype = dtype or first.meta['dtype']
-                kwargs['dtype'] = dtype
+            first = sources[0]
+            kwargs = first.profile
+            kwargs.update(**creation_options)
+            dtype = dtype or first.meta['dtype']
+            kwargs['dtype'] = dtype
+
+            # Extend snuggs.
+            snuggs.func_map['read'] = _read_array
+            snuggs.func_map['band'] = lambda d, i: _get_bands(inputs, sources, d, i)
+            snuggs.func_map['bands'] = lambda d: _get_bands(inputs, sources, d)
+            snuggs.func_map['fillnodata'] = lambda *args: fillnodata(*args)
+            snuggs.func_map['sieve'] = lambda *args: sieve(*args)
+
+            dst = None
+
+            # The windows iterator is initialized with a single sample.
+            # The actual work windows will be added in the second
+            # iteration of the loop.
+            work_windows = [(None, Window(0, 0, 16, 16))]
+
+            for ij, window in work_windows:
+
+                ctxkwds = OrderedDict()
+
+                for i, ((name, path), src) in enumerate(zip(inputs, sources)):
 
-            ctxkwds = OrderedDict()
-            for i, (name, path) in enumerate(inputs):
-                with rasterio.open(path) as src:
                     # Using the class method instead of instance
                     # method. Latter raises
                     #
@@ -110,32 +177,31 @@ def calc(ctx, command, files, output, name, dtype, masked, overwrite,
                     #
                     # possibly something to do with the instance being
                     # a masked array.
-                    ctxkwds[name or '_i%d' % (i + 1)] = src.read(masked=masked)
-
-            # Extend snuggs.
-            snuggs.func_map['read'] = read_array
-            snuggs.func_map['band'] = lambda d, i: get_bands(inputs, d, i)
-            snuggs.func_map['bands'] = lambda d: get_bands(inputs, d)
-            snuggs.func_map['fillnodata'] = lambda *args: fillnodata(*args)
-            snuggs.func_map['sieve'] = lambda *args: sieve(*args)
+                    ctxkwds[name or '_i%d' % (i + 1)] = src.read(masked=masked, window=window)
 
-            res = snuggs.eval(command, ctxkwds)
+                res = snuggs.eval(command, **ctxkwds)
 
-            if (isinstance(res, np.ma.core.MaskedArray) and (
-                    tuple(LooseVersion(np.__version__).version) < (1, 9) or
-                    tuple(LooseVersion(np.__version__).version) > (1, 10))):
-                res = res.filled(kwargs['nodata'])
+                if (isinstance(res, np.ma.core.MaskedArray) and (
+                        tuple(LooseVersion(np.__version__).version) < (1, 9) or
+                        tuple(LooseVersion(np.__version__).version) > (1, 10))):
+                    res = res.filled(kwargs['nodata'])
 
-            if len(res.shape) == 3:
-                results = np.ndarray.astype(res, dtype, copy=False)
-            else:
-                results = np.asanyarray(
-                    [np.ndarray.astype(res, dtype, copy=False)])
+                if len(res.shape) == 3:
+                    results = np.ndarray.astype(res, dtype, copy=False)
+                else:
+                    results = np.asanyarray(
+                        [np.ndarray.astype(res, dtype, copy=False)])
 
-            kwargs['count'] = results.shape[0]
+                # The first iteration is only to get sample results and from them
+                # compute some properties of the output dataset.
+                if dst is None:
+                    kwargs['count'] = results.shape[0]
+                    dst = rasterio.open(output, 'w', **kwargs)
+                    work_windows.extend(_chunk_output(dst.width, dst.height, dst.count, np.dtype(dst.dtypes[0]).itemsize, mem_limit=mem_limit))
 
-            with rasterio.open(output, 'w', **kwargs) as dst:
-                dst.write(results)
+                # In subsequent iterations we write results.
+                else:
+                    dst.write(results, window=window)
 
     except snuggs.ExpressionError as err:
         click.echo("Expression Error:")
@@ -143,3 +209,9 @@ def calc(ctx, command, files, output, name, dtype, masked, overwrite,
         click.echo(' ' + ' ' * err.offset + "^")
         click.echo(err)
         raise click.Abort()
+
+    finally:
+        if dst:
+            dst.close()
+        for src in sources:
+            src.close()


Test patch:
diff --git a/tests/test_rio_calc.py b/tests/test_rio_calc.py
--- a/tests/test_rio_calc.py
+++ b/tests/test_rio_calc.py
@@ -1,15 +1,11 @@
-import sys
-import logging
-
 from click.testing import CliRunner
+import pytest
 
 import rasterio
+from rasterio.rio.calc import _chunk_output
 from rasterio.rio.main import main_group
 
 
-logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
-
-
 def test_err(tmpdir):
     outfile = str(tmpdir.join('out.tif'))
     runner = CliRunner()
@@ -138,6 +134,7 @@ def test_fillnodata_map(tmpdir):
         assert round(data.mean(), 1) == 58.6
         assert data[0][60][60] > 0
 
+
 def test_sieve_band(tmpdir):
     outfile = str(tmpdir.join('out.tif'))
     runner = CliRunner()
@@ -183,3 +180,16 @@ def test_positional_calculation_byindex(tmpdir):
 
     with rasterio.open(outfile) as src:
         assert src.read(1, window=window) == answer
+
+
+@pytest.mark.parametrize('width', [10, 791, 3000])
+@pytest.mark.parametrize('height', [8, 718, 4000])
+@pytest.mark.parametrize('count', [1, 3, 4])
+@pytest.mark.parametrize('itemsize', [1, 2, 8])
+@pytest.mark.parametrize('mem_limit', [1, 16, 64, 512])
+def test_chunk_output(width, height, count, itemsize, mem_limit):
+    work_windows = _chunk_output(width, height, count, itemsize, mem_limit=mem_limit)
+    num_windows_rows = max([i for ((i, j), w) in work_windows]) + 1
+    num_windows_cols = max([j for ((i, j), w) in work_windows]) + 1
+    assert sum((w.width for ij, w in work_windows)) == width * num_windows_rows
+    assert sum((w.height for ij, w in work_windows)) == height * num_windows_cols