Repository: dcs4cop/xcube
Pull request #160 (2019-08-21 14:08:05)

Forman new cube with var funcs
`xcube.api.new_cube` function now accepts callables as values for variables.
This allows to compute variable values depending on the (t, y, x) position
in the cube. Useful for testing. 

I need PR accepted for another PR regarding replacement of time slices.


Patch:
diff --git a/CHANGES.md b/CHANGES.md
index d6e3538d..3c7f3d0a 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -2,6 +2,10 @@
 
 ### Enhancements
 
+* `xcube.api.new_cube` function now accepts callables as values for variables.
+  This allows to compute variable values depending on the (t, y, x) position
+  in the cube. Useful for testing.
+* Added new CLI tool `xcube optimize` and API function `xcube.api.optimize_dataset`
 * `xcube.api` now exports the `MaskSet` class which is useful for decoding flag values encoding following the
   [CF conventions](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.7/cf-conventions.html#flags).
 * Added new CLI tool `xcube optimize` and API function `xcube.api.optimize_dataset` 
diff --git a/xcube/api/new.py b/xcube/api/new.py
index d0ee3f1b..8e58db82 100644
--- a/xcube/api/new.py
+++ b/xcube/api/new.py
@@ -1,3 +1,5 @@
+import itertools
+
 import numpy as np
 import pandas as pd
 import xarray as xr
@@ -20,7 +22,15 @@ def new_cube(title="Test Cube",
              drop_bounds=False,
              variables=None):
     """
-    Create a new empty cube. Useful for testing.
+    Create a new empty cube. Useful for creating cubes templates with predefined
+    coordinate variables and metadata. The function is also heavily used
+    by xcube's unit tests.
+
+    The values of the *variables* dictionary can be either constants,
+    array-like objects, or functions that compute their return value from
+    passed coordinate indexes. The expected signature is:::
+
+        def my_func(time: int, lat: int, lon: int) -> Union[bool, int, float]
 
     :param title: A title.
     :param width: Horizontal number of grid cells
@@ -109,8 +119,6 @@ def new_cube(title="Test Cube",
         "geospatial_lat_units": "degrees_north",
     }
 
-    # TODO (forman): allow variable values to be expressions so values will be computed from coords using numexpr
-
     data_vars = {}
     if variables:
         dims = ("time", "lat", "lon")
@@ -121,6 +129,12 @@ def new_cube(title="Test Cube",
                 data_vars[var_name] = data
             elif isinstance(data, int) or isinstance(data, float) or isinstance(data, bool):
                 data_vars[var_name] = xr.DataArray(np.full(shape, data), dims=dims)
+            elif callable(data):
+                func = data
+                data = np.zeros(shape)
+                for index in itertools.product(*map(range, shape)):
+                    data[index] = func(*index)
+                data_vars[var_name] = xr.DataArray(data, dims=dims)
             elif data is None:
                 data_vars[var_name] = xr.DataArray(np.random.uniform(0.0, 1.0, size).reshape(shape), dims=dims)
             else:


Test patch:
diff --git a/test/api/test_new.py b/test/api/test_new.py
index 8adb4fa9..0764acea 100644
--- a/test/api/test_new.py
+++ b/test/api/test_new.py
@@ -1,5 +1,7 @@
 import unittest
 
+import numpy as np
+
 from xcube.api import new_cube
 
 
@@ -24,3 +26,39 @@ class NewCubeTest(unittest.TestCase):
     def test_new_cube_without_bounds(self):
         cube = new_cube(drop_bounds=True)
         self.assertEqual({'lon': 360, 'lat': 180, 'time': 5}, cube.dims)
+
+    def test_new_cube_with_const_vars(self):
+        cube = new_cube(variables=dict(sst=274.4, chl=10.31))
+        self.assertIn('sst', cube)
+        self.assertIn('chl', cube)
+        import numpy as np
+        np.testing.assert_almost_equal(cube.sst.values, np.full((5, 180, 360), 274.4))
+        np.testing.assert_almost_equal(cube.chl.values, np.full((5, 180, 360), 10.31))
+
+    def test_new_cube_with_func_vars(self):
+        def sst_func(t, y, x):
+            return 274.4 + t * 0.5
+
+        def chl_func(t, y, x):
+            return 10.31 + y * 0.1
+
+        def aot_func(t, y, x):
+            return 0.88 + x * 0.05
+
+        cube = new_cube(width=16, height=8, variables=dict(sst=sst_func, chl=chl_func, aot=aot_func))
+
+        self.assertIn('sst', cube)
+        self.assertIn('chl', cube)
+        self.assertIn('aot', cube)
+
+        np.testing.assert_almost_equal(cube.sst[0, :, :].values, np.full((8, 16), 274.4))
+        np.testing.assert_almost_equal(cube.chl[:, 0, :].values, np.full((5, 16), 10.31))
+        np.testing.assert_almost_equal(cube.aot[:, :, 0].values, np.full((5, 8), 0.88))
+
+        np.testing.assert_almost_equal(cube.sst.isel(lon=3, lat=2).values,
+                                       np.array([274.4, 274.9, 275.4, 275.9, 276.4]))
+        np.testing.assert_almost_equal(cube.chl.isel(time=3, lon=4).values,
+                                       np.array([10.31, 10.41, 10.51, 10.61, 10.71, 10.81, 10.91, 11.01]))
+        np.testing.assert_almost_equal(cube.aot.isel(time=1, lat=2).values,
+                                       np.array([0.88, 0.93, 0.98, 1.03, 1.08, 1.13, 1.18, 1.23,
+                                                 1.28, 1.33, 1.38, 1.43, 1.48, 1.53, 1.58, 1.63]))