Repository: dcs4cop/xcube
Pull request #258 (2020-01-14 16:36:16)

Support xarray and numpy functions in expressions
Now supporting xarray and numpy functions in expressions. You can now use 
`xr` and `np` contexts, e.g. `xr.where(CHL >= 0.0, CHL)` or `CHL.where(CHL >= 0.0)`.

Closes  #257


Patch:
diff --git a/CHANGES.md b/CHANGES.md
index 4ef8069d..ef104c8f 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,9 +1,13 @@
 ## Changes in 0.4.0 (in development)
 
+### New
+
+* Now supporting xarray and numpy functions in expressions. You can now use 
+  `xr` and `np` contexts, e.g. `xr.where(CHL >= 0.0, CHL)`. (#257)
+
 ### Other
 
 * Renamed default log file for `xcube serve` command to `xcube-serve.log`.
-* `xcube gen` now immediately flushes logging output to standard out
   
 ## Changes in 0.3.1 (in development)
 
diff --git a/xcube/cli/gen.py b/xcube/cli/gen.py
index 1b5f50de..001c2bd3 100644
--- a/xcube/cli/gen.py
+++ b/xcube/cli/gen.py
@@ -123,11 +123,8 @@ def gen(input: Sequence[str],
         no_sort_mode=no_sort,
     )
 
-    def flushing_monitor(*args, **kwargs):
-        print(*args, flush=True, **kwargs)
-
     gen_cube(dry_run=dry_run,
-             monitor=flushing_monitor,
+             monitor=print,
              **config)
 
     return 0
diff --git a/xcube/core/gen/gen.py b/xcube/core/gen/gen.py
index 8e98bc00..41c1f801 100644
--- a/xcube/core/gen/gen.py
+++ b/xcube/core/gen/gen.py
@@ -206,7 +206,7 @@ def _process_input(input_processor: InputProcessor,
         return False
 
     if output_variables:
-        output_variables = to_resolved_name_dict_pairs(output_variables, input_dataset)
+        output_variables = to_resolved_name_dict_pairs(output_variables, input_dataset, keep=True)
     else:
         output_variables = [(var_name, None) for var_name in input_dataset.data_vars]
 
diff --git a/xcube/util/expression.py b/xcube/util/expression.py
index a95ad60a..cc2887c1 100644
--- a/xcube/util/expression.py
+++ b/xcube/util/expression.py
@@ -62,7 +62,7 @@ def compute_expr(expr: str,
 
 def transpile_expr(expr: str, warn=False) -> str:
     """
-    Transpile a Python expression into a numpy array expression.
+    Transpile a Python expression into a numpy/xarray array expression.
 
     :param expr: The Python expression
     :param warn: If warnings shall be emitted
@@ -74,7 +74,7 @@ def transpile_expr(expr: str, warn=False) -> str:
 # noinspection PyMethodMayBeStatic
 class _ExprTranspiler:
     """
-    Transpiles a BEAM/SNAP expression into a numpy array expression.
+    Transpile a Python expression into a numpy/xarray array expression.
 
     See https://greentreesnakes.readthedocs.io/en/latest/nodes.html#expressions
     """
@@ -182,8 +182,14 @@ class _ExprTranspiler:
         return "%s(%s)" % (self.transform_function_name(func), args)
 
     def transform_function_name(self, func: ast.Name):
-        name = dict(min='fmin', max='fmax').get(func.id, func.id)
-        return 'np.%s' % name
+        if isinstance(func, ast.Attribute):
+            return '%s.%s' % (func.value.id, func.attr)
+        else:
+            if func.id in ('where',):
+                return 'xr.%s' % func.id
+            else:
+                name = dict(min='fmin', max='fmax').get(func.id, func.id)
+                return 'np.%s' % name
 
     # noinspection PyUnusedLocal
     def transform_attribute(self, value: ast.AST, attr: str, ctx):
@@ -237,7 +243,7 @@ class _ExprTranspiler:
 
     # noinspection PyUnusedLocal
     def transform_if_exp(self, test, body, orelse):
-        return 'np.where({y}, {x}, {z})'
+        return 'xr.where({y}, {x}, {z})'
 
     def transform_bool_op(self, op, values):
         name, precedence, assoc = _ExprTranspiler.get_op_info(op)


Test patch:
diff --git a/test/util/test_expression.py b/test/util/test_expression.py
index 997daa22..41f2bcc5 100644
--- a/test/util/test_expression.py
+++ b/test/util/test_expression.py
@@ -1,6 +1,7 @@
 import unittest
 
 import numpy as np
+import xarray as xr
 
 from xcube.util.expression import transpile_expr, compute_expr, compute_array_expr
 
@@ -10,7 +11,8 @@ class ComputeArrayExprTest(unittest.TestCase):
     def test_valid_exprs(self):
         namespace = dict(a=np.array([0.1, 0.3, 0.1, 0.7, 0.4, 0.9]),
                          b=np.array([0.2, 0.1, 0.3, 0.2, 0.4, 0.8]),
-                         np=np)
+                         np=np,
+                         xr=xr)
 
         value = compute_array_expr('a + 1', namespace=namespace)
         np.testing.assert_array_almost_equal(value,
@@ -32,6 +34,16 @@ class ComputeArrayExprTest(unittest.TestCase):
         np.testing.assert_equal(value,
                                 np.array([False, False, False, False, True, False]))
 
+        # This weirdo expression is a result of translating SNAP conditional expressions to Python.
+        value = compute_array_expr('a > 0.35 if a else b', namespace=namespace)
+        np.testing.assert_equal(value,
+                                np.array([0.2, 0.1, 0.3, 0.7, 0.4, 0.9]))
+
+        # We actually mean
+        value = compute_array_expr('where(a > 0.35, a, b)', namespace=namespace)
+        np.testing.assert_equal(value,
+                                np.array([0.2, 0.1, 0.3, 0.7, 0.4, 0.9]))
+
     def test_invalid_exprs(self):
         with self.assertRaises(ValueError) as cm:
             compute_expr('20. -')
@@ -148,9 +160,20 @@ class TranspileExprTest(unittest.TestCase):
         # The following conditional expr looks wrong but this is how it looks like after translating
         # from SNAP expression 'a >= 0.0 ? a : NaN'
         self.assertEqual(transpile_expr('a >= 0.0 if a else NaN'),
-                         'np.where(a >= 0.0, a, NaN)')
+                         'xr.where(a >= 0.0, a, NaN)')
         self.assertEqual(transpile_expr('a >= 0.0 if a else b >= 0.0 if b else NaN'),
-                         'np.where(a >= 0.0, a, np.where(b >= 0.0, b, NaN))')
+                         'xr.where(a >= 0.0, a, xr.where(b >= 0.0, b, NaN))')
+
+    def test_where(self):
+        self.assertEqual(transpile_expr('where(a >= 0.0, a, NaN)'),
+                         'xr.where(a >= 0.0, a, NaN)')
+        self.assertEqual(transpile_expr('xr.where(a >= 0.0, a, NaN)'),
+                         'xr.where(a >= 0.0, a, NaN)')
+        self.assertEqual(transpile_expr('np.where(a >= 0.0, a, NaN)'),
+                         'np.where(a >= 0.0, a, NaN)')
+        # xarray.DataArray.where() method:
+        self.assertEqual(transpile_expr('a.where(a.x >= 0.0)'),
+                         'a.where(a.x >= 0.0)')
 
     def test_mixed(self):
         self.assertEqual(transpile_expr('a+sin(x + 2.8)'),