Repository: opendatacube/datacube-core

Issue:
Need formula to respect `nodata` in `virtual product`
### Expected behaviour
Formula should respect the `nodata` defined in the input and output


Discussion:

Resolving patch:
diff --git a/datacube/storage/masking.py b/datacube/storage/masking.py
--- a/datacube/storage/masking.py
+++ b/datacube/storage/masking.py
@@ -119,7 +119,8 @@ def valid_data_mask(data):
 
     nodata = data.attrs.get('nodata', None)
 
-    return xarray.apply_ufunc(lambda xx: valid_mask(xx, nodata), data)
+    return xarray.apply_ufunc(lambda xx: valid_mask(xx, nodata), data, dask='parallelized',
+                              output_dtypes=[numpy.bool])
 
 
 def mask_valid_data(data, keep_attrs=True):
diff --git a/datacube/virtual/__init__.py b/datacube/virtual/__init__.py
--- a/datacube/virtual/__init__.py
+++ b/datacube/virtual/__init__.py
@@ -2,7 +2,7 @@
 
 from .impl import VirtualProduct, Transformation, VirtualProductException
 from .impl import from_validated_recipe
-from .transformations import MakeMask, ApplyMask, ToFloat, Rename, Select, Formula
+from .transformations import MakeMask, ApplyMask, ToFloat, Rename, Select, Expressions
 from .transformations import Mean, year, month, week, day
 from .catalog import Catalog
 from .utils import reject_keys
@@ -112,7 +112,7 @@ def register(self, namespace: str, name: str, callable_obj):
                                                    to_float=ToFloat,
                                                    rename=Rename,
                                                    select=Select,
-                                                   formula=Formula),
+                                                   expressions=Expressions),
                                  'aggregate': dict(mean=Mean),
                                  'aggregate/group_by': dict(year=year,
                                                             month=month,
diff --git a/datacube/virtual/transformations.py b/datacube/virtual/transformations.py
--- a/datacube/virtual/transformations.py
+++ b/datacube/virtual/transformations.py
@@ -4,6 +4,9 @@
 
 from datacube.storage.masking import make_mask as make_mask_prim
 from datacube.storage.masking import mask_invalid_data as mask_invalid_data_prim
+from datacube.storage.masking import valid_data_mask
+
+from datacube.utils.math import dtype_is_float
 
 from .impl import VirtualProductException, Transformation, Measurement
 
@@ -84,6 +87,7 @@ def worker(_, value):
 
             result = value.copy()
             result['dtype'] = self.fallback_dtype
+            result['nodata'] = float('nan')
             return Measurement(**result)
 
         return selective_apply_dict(rest, apply_to=self.apply_to, value_map=worker)
@@ -110,7 +114,9 @@ def worker(key, value):
                     raise VirtualProductException("measurement {} has no nodata value".format(key))
                 return value.where(mask, value.attrs['nodata'])
 
-            return value.where(mask).astype(self.fallback_dtype)
+            result = value.where(mask).astype(self.fallback_dtype)
+            result.attrs['nodata'] = float('nan')
+            return result
 
         return selective_apply(rest, apply_to=self.apply_to, value_map=worker)
 
@@ -244,9 +250,10 @@ class EvaluateTree(lark.Transformer):
     int_literal = int
 
 
-class Formula(Transformation):
-    def __init__(self, output):
+class Expressions(Transformation):
+    def __init__(self, output, masked=True):
         self.output = output
+        self.masked = masked
 
     def measurements(self, input_measurements):
         parser = formula_parser()
@@ -267,7 +274,8 @@ def deduce_type(output_var, output_desc):
 
         def measurement(output_var, output_desc):
             return Measurement(name=output_var, dtype=deduce_type(output_var, output_desc),
-                               nodata=output_desc.get('nodata'), units=output_desc.get('units'))
+                               nodata=output_desc.get('nodata', float('nan')),
+                               units=output_desc.get('units', '1'))
 
         return {output_var: measurement(output_var, output_desc)
                 for output_var, output_desc in self.output.items()}
@@ -280,13 +288,67 @@ class EvaluateData(EvaluateTree):
             def var_name(self, key):
                 return data[key.value]
 
-        ev = EvaluateData()
+        @lark.v_args(inline=True)
+        class EvaluateNodataMask(lark.Transformer):
+            # the result of an expression is nodata whenever any of its subexpressions is nodata
+            from operator import or_
+
+            # pylint: disable=invalid-name
+            and_ = _xor = or_
+            eq = ne = le = ge = lt = gt = or_
+            add = sub = mul = truediv = floordiv = mod = pow = lshift = rshift = or_
+
+            def not_(self, value):
+                return value
+
+            neg = pos = inv = not_
+
+            @staticmethod
+            def float_literal(value):
+                return False
+
+            @staticmethod
+            def int_literal(value):
+                return False
+
+            def var_name(self, key):
+                # pylint: disable=invalid-unary-operand-type
+                return ~valid_data_mask(data[key.value])
+
+        ev_data = EvaluateData()
+        ev_mask = EvaluateNodataMask()
 
         def result(output_var, output_desc):
+            # pylint: disable=invalid-unary-operand-type
+
+            nodata = output_desc.get('nodata')
+
             formula = output_desc['formula']
             tree = parser.parse(formula)
-
-            return ev.transform(tree)
+            result = ev_data.transform(tree)
+            result.attrs['crs'] = data.attrs['crs']
+            if nodata is not None:
+                result.attrs['nodata'] = nodata
+            result.attrs['units'] = output_desc.get('units', '1')
+
+            if not self.masked:
+                return result
+
+            # masked output
+            dtype = result.dtype
+            mask = ev_mask.transform(tree)
+
+            if not dtype_is_float(dtype) and nodata is None:
+                raise VirtualProductException("cannot mask without specified nodata")
+
+            if nodata is None:
+                result = result.where(~mask)
+                result.attrs['nodata'] = numpy.nan
+            else:
+                result = result.where(~mask, nodata)
+                result.attrs['nodata'] = nodata
+
+            return result
 
         return xarray.Dataset(data_vars={output_var: result(output_var, output_desc)
                                          for output_var, output_desc in self.output.items()},


Test patch:
diff --git a/tests/api/test_virtual.py b/tests/api/test_virtual.py
--- a/tests/api/test_virtual.py
+++ b/tests/api/test_virtual.py
@@ -345,12 +345,13 @@ def test_aliases(dc, query):
     assert 'green' not in data
 
 
-def test_formula(dc, query):
+def test_expressions(dc, query):
     bluegreen = construct_from_yaml("""
-        transform: formula
+        transform: expressions
         output:
             bluegreen:
                 formula: blue + green
+                nodata: -999
         input:
             product: ls8_nbar_albers
             measurements: [blue, green]
@@ -366,6 +367,7 @@ def test_formula(dc, query):
         data = bluegreen.load(dc, **query)
 
     assert 'bluegreen' in data
+    assert numpy.all((data.bluegreen == -999).values)
 
 
 def test_aggregate(dc, query, catalog):