From 930c63f84262e9124bd0d488d67efd0449634e8f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dion=20H=C3=A4fner?= <dion.haefner@nbi.ku.dk>
Date: Thu, 8 Jul 2021 10:53:32 +0200
Subject: [PATCH] add transparency mask support for /compute

---
 terracotta/expressions.py      | 25 ++++++++++++++--
 tests/handlers/test_compute.py | 52 ++++++++++++++++++++++++++++++++++
 tests/test_expressions.py      | 19 +++++++++++++
 3 files changed, 93 insertions(+), 3 deletions(-)

diff --git a/terracotta/expressions.py b/terracotta/expressions.py
index 3169ef89..103cf083 100644
--- a/terracotta/expressions.py
+++ b/terracotta/expressions.py
@@ -12,8 +12,22 @@
 
 
 EXTRA_CALLABLES = {
-    # name: (callable, nargs)
-    'where': (np.where, 3),
+    # 'name': (callable, nargs)
+
+    # boolean ops
+    'where': (np.ma.where, 3),
+    'masked_equal': (np.ma.masked_equal, 2),
+    'masked_greater': (np.ma.masked_greater, 2),
+    'masked_greater_equal': (np.ma.masked_greater_equal, 2),
+    'masked_inside': (np.ma.masked_inside, 3),
+    'masked_invalid': (np.ma.masked_invalid, 1),
+    'masked_less': (np.ma.masked_less, 2),
+    'masked_less_equal': (np.ma.masked_less_equal, 2),
+    'masked_not_equal': (np.ma.masked_not_equal, 2),
+    'masked_outside': (np.ma.masked_outside, 3),
+    'masked_where': (np.ma.masked_where, 2),
+
+    # math
     'minimum': (np.minimum, 2),
     'maximum': (np.maximum, 2),
     'abs': (np.abs, 1),
@@ -37,7 +51,9 @@
 }
 
 EXTRA_CONSTANTS = {
-    'pi': np.pi
+    'pi': np.pi,
+    'nan': np.nan,
+    'inf': np.inf,
 }
 
 
@@ -173,4 +189,7 @@ def evaluate_expression(expr: str,
     if not isinstance(result, np.ndarray):
         raise ValueError('expression does not return an array')
 
+    # mask inf and nan values
+    result = np.ma.masked_invalid(result)
+
     return result
diff --git a/tests/handlers/test_compute.py b/tests/handlers/test_compute.py
index 1503ac30..2b341126 100644
--- a/tests/handlers/test_compute.py
+++ b/tests/handlers/test_compute.py
@@ -46,3 +46,55 @@ def test_compute_consistency(use_testdb, testdb, raster_file_xyz):
         img_data,
         to_uint8(v1 + v2, 0, 10000)
     )
+
+
+def test_compute_transparency_nan(use_testdb, testdb, raster_file_xyz):
+    import terracotta
+    from terracotta.xyz import get_tile_data
+    from terracotta.handlers import compute
+
+    raw_img = compute.compute(
+        'where(v1 > 0, nan, v1 + v2)',
+        ['val21', 'x'],
+        {'v1': 'val22', 'v2': 'val23'},
+        stretch_range=(0, 10000),
+        tile_xyz=raster_file_xyz
+    )
+    img_data = np.asarray(Image.open(raw_img).convert('RGBA'))
+
+    driver = terracotta.get_driver(testdb)
+
+    with driver.connect():
+        v1 = get_tile_data(driver, ['val21', 'x', 'val22'], raster_file_xyz)
+
+    alpha = img_data[..., 3]
+    np.testing.assert_array_equal(
+        alpha,
+        np.where(v1.mask | (v1 > 0), 0, 255),
+    )
+
+
+def test_compute_transparency_mask(use_testdb, testdb, raster_file_xyz):
+    import terracotta
+    from terracotta.xyz import get_tile_data
+    from terracotta.handlers import compute
+
+    raw_img = compute.compute(
+        'masked_where(v1 > 0, v1 + v2)',
+        ['val21', 'x'],
+        {'v1': 'val22', 'v2': 'val23'},
+        stretch_range=(0, 10000),
+        tile_xyz=raster_file_xyz
+    )
+    img_data = np.asarray(Image.open(raw_img).convert('RGBA'))
+
+    driver = terracotta.get_driver(testdb)
+
+    with driver.connect():
+        v1 = get_tile_data(driver, ['val21', 'x', 'val22'], raster_file_xyz)
+
+    alpha = img_data[..., 3]
+    np.testing.assert_array_equal(
+        alpha,
+        np.where(v1.mask | (v1 > 0), 0, 255),
+    )
diff --git a/tests/test_expressions.py b/tests/test_expressions.py
index 24f28a4e..1f6e3fe4 100644
--- a/tests/test_expressions.py
+++ b/tests/test_expressions.py
@@ -226,3 +226,22 @@ def test_timeout():
         evaluate_expression('+'.join(['v1'] * 10), {'v1': np.ones((256, 256))}, timeout=0)
 
     assert 'timeout' in str(raised_exc.value)
+
+
+def test_mask_invalid():
+    from terracotta.expressions import evaluate_expression
+    res = evaluate_expression('where(v1 + v2 < 1, nan, 0)', OPERANDS)
+    mask = OPERANDS['v1'] + OPERANDS['v2'] < 1
+
+    assert isinstance(res, np.ma.MaskedArray)
+    assert np.all(res == 0)
+    assert np.array_equal(res.mask, mask)
+
+
+def test_out_dtype():
+    from terracotta.expressions import evaluate_expression
+    operands = dict(v1=np.ones(10, dtype='int64'), v2=np.zeros(10, dtype='int32'))
+    res = evaluate_expression('v1 + v2', operands)
+
+    assert isinstance(res, np.ma.MaskedArray)
+    assert res.dtype == np.dtype('int64')