Merge 82dec1c into 3f4d5ed

rasterio/rio/calc.py

@@ -0,0 +1,121 @@
+# Calc command.
+
+import logging
+import math
+import os.path
+import re
+import sys
+import warnings
+
+import click
+import snuggs
+
+import rasterio
+from rasterio.rio.cli import cli
+
+
+@cli.command(short_help="Raster data calculator.")
+@click.argument('command')
+@click.argument(
+    'files',
+    nargs=-1,
+    type=click.Path(resolve_path=False),
+    required=True,
+    metavar="INPUTS... OUTPUT")
+@click.option('--name', multiple=True,
+        help='Specify an input file with a unique short (alphas only) name '
+             'for use in commands like "a=tests/data/RGB.byte.tif".')
+@click.option('--dtype', 
+              type=click.Choice([
+                'ubyte', 'uint8', 'uint16', 'int16', 'uint32',
+                'int32', 'float32', 'float64']),
+                default='float64',
+              help="Output data type (default: float64).")
+@click.pass_context
+def calc(ctx, command, files, name, dtype):
+    """A raster data calculator
+
+    Evaluates an expression using input datasets and writes the result
+    to a new dataset.
+
+    Command syntax is lisp-like. An expression consists of an operator
+    or function name and one or more strings, numbers, or expressions
+    enclosed in parentheses. Functions include ``read`` (gets a raster
+    array) and ``asarray`` (makes a 3-D array from 2-D arrays).
+
+    \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).
+        * 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 single array, a single band output
+          file is written.
+
+    Example:
+
+    \b
+         $ rio calc "(+ 2 (* 0.95 (read 1)))" tests/data/RGB.byte.tif \\
+         > /tmp/out.tif --dtype ubyte
+
+    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 --dtype ubyte
+
+    Produces a 3-band RGB GeoTIFF, with red levels incremented by 125,
+    from the single-band input.
+
+    """
+    import numpy as np
+
+    verbosity = (ctx.obj and ctx.obj.get('verbosity')) or 1
+    logger = logging.getLogger('rio')
+
+    try:
+        with rasterio.drivers(CPL_DEBUG=verbosity>2):
+            output = files[-1]
+
+            inputs = (
+                    [tuple(n.split('=')) for n in name] +
+                    [(None, n) for n in files[:-1]])
+
+            with rasterio.open(inputs[0][1]) as first:
+                kwargs = first.meta
+                kwargs['transform'] = kwargs.pop('affine')
+                kwargs['dtype'] = dtype
+
+            ctxkwds = {}
+            for i, (name, path) in enumerate(inputs):
+                with rasterio.open(path) as src:
+                    # Using the class method instead of instance method.
+                    # Latter raises
+                    #
+                    # TypeError: astype() got an unexpected keyword argument 'copy'
+                    # 
+                    # possibly something to do with the instance being a masked
+                    # array.
+                    ctxkwds[name or '_i%d' % (i+1)] = np.ndarray.astype(
+                            src.read(), 'float64', copy=False)
+
+            res = snuggs.eval(command, **ctxkwds)
+
+            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]
+            with rasterio.open(output, 'w', **kwargs) as dst:
+                dst.write(results)
+
+        sys.exit(0)
+    except Exception:
+        logger.exception("Failed. Exception caught")
+        sys.exit(1)

rasterio/rio/main.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 
+from rasterio.rio.calc import calc
 from rasterio.rio.cli import cli
 from rasterio.rio.bands import stack
 from rasterio.rio.features import shapes, rasterize

requirements-dev.txt

@@ -5,6 +5,7 @@ cython==0.21.2
 delocate
 enum34
 numpy>=1.8.0
+snuggs>=1.0
 pytest
 setuptools>=0.9.8
 wheel

requirements.txt

@@ -2,4 +2,5 @@ affine
 cligj
 enum34
 numpy>=1.8.0
+snuggs>=1.0
 setuptools

setup.py

@@ -169,7 +169,8 @@
 inst_reqs = [
     'affine>=1.0',
     'cligj',
-    'Numpy>=1.7' ]
+    'Numpy>=1.7',
+    'snuggs']
 
 if sys.version_info < (3, 4):
     inst_reqs.append('enum34')

tests/test_rio_calc.py

@@ -0,0 +1,121 @@
+import sys
+import logging
+
+from click.testing import CliRunner
+
+import rasterio
+from rasterio.rio.calc import calc
+
+
+logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
+
+
+def test_err(tmpdir):
+    outfile = str(tmpdir.join('out.tif'))
+    runner = CliRunner()
+    result = runner.invoke(calc, [
+                '($ 0.1 (read 1))', 'tests/data/shade.tif', outfile],
+                catch_exceptions=False)
+    assert result.exit_code == 1
+
+
+def test_multiband_calc(tmpdir):
+    outfile = str(tmpdir.join('out.tif'))
+    runner = CliRunner()
+    result = runner.invoke(calc, [
+                '(+ 125 (* 0.1 (read 1)))', 'tests/data/shade.tif', outfile],
+                catch_exceptions=False)
+    assert result.exit_code == 0
+    with rasterio.open(outfile) as src:
+        assert src.count == 1
+        assert src.meta['dtype'] == 'float64'
+        data = src.read()
+        assert data.min() == 125
+
+
+def test_singleband_calc_byindex(tmpdir):
+    outfile = str(tmpdir.join('out.tif'))
+    runner = CliRunner()
+    result = runner.invoke(calc, [
+                    '(+ 125 (* 0.1 (read 1 1)))',
+                    'tests/data/shade.tif',
+                    outfile],
+                catch_exceptions=False)
+    assert result.exit_code == 0
+    with rasterio.open(outfile) as src:
+        assert src.count == 1
+        assert src.meta['dtype'] == 'float64'
+        data = src.read()
+        assert data.min() == 125
+
+
+def test_singleband_calc_byname(tmpdir):
+    outfile = str(tmpdir.join('out.tif'))
+    runner = CliRunner()
+    result = runner.invoke(calc, [
+                    '(+ 125 (* 0.1 (take shade 1)))',
+                    '--name', 'shade=tests/data/shade.tif',
+                    outfile],
+                catch_exceptions=False)
+    assert result.exit_code == 0
+    with rasterio.open(outfile) as src:
+        assert src.count == 1
+        assert src.meta['dtype'] == 'float64'
+        data = src.read()
+        assert data.min() == 125
+
+
+def test_parts_calc(tmpdir):
+    # Producing an RGB output from the hill shade.
+    # Red band has bumped up values. Other bands are unchanged.
+    outfile = str(tmpdir.join('out.tif'))
+    runner = CliRunner()
+    result = runner.invoke(calc, [
+                    '(asarray (+ (read 1 1) 125) (read 1 1) (read 1 1))',
+                    '--dtype', 'uint8',
+                    'tests/data/shade.tif',
+                    outfile],
+                catch_exceptions=False)
+    assert result.exit_code == 0
+    with rasterio.open(outfile) as src:
+        assert src.count == 3
+        assert src.meta['dtype'] == 'uint8'
+        data = src.read()
+        assert data[0].min() == 125
+        assert data[1].min() == 0
+        assert data[2].min() == 0
+
+
+def test_parts_calc_2(tmpdir):
+    # Produce greyscale output from the RGB file.
+    outfile = str(tmpdir.join('out.tif'))
+    runner = CliRunner()
+    result = runner.invoke(calc, [
+                    '(+ (+ (/ (read 1 1) 3) (/ (read 1 2) 3)) (/ (read 1 3) 3))',
+                    '--dtype', 'uint8',
+                    'tests/data/RGB.byte.tif',
+                    outfile],
+                catch_exceptions=False)
+    assert result.exit_code == 0
+    with rasterio.open(outfile) as src:
+        assert src.count == 1
+        assert src.meta['dtype'] == 'uint8'
+        data = src.read()
+        assert round(data.mean(), 1) == 60.3
+
+
+def test_copy_rgb(tmpdir):
+    outfile = str(tmpdir.join('out.tif'))
+    runner = CliRunner()
+    result = runner.invoke(calc, [
+                    '(read 1)',
+                    '--dtype', 'uint8',
+                    'tests/data/RGB.byte.tif',
+                    outfile],
+                catch_exceptions=False)
+    assert result.exit_code == 0
+    with rasterio.open(outfile) as src:
+        assert src.count == 3
+        assert src.meta['dtype'] == 'uint8'
+        data = src.read()
+        assert round(data.mean(), 1) == 60.6