Repository: dcs4cop/xcube
Pull request #273 (2020-02-06 09:44:45)

Forman 272 computed vars
## What's in here

The configuration of `xcube serve` has been enhanced to support augmentation of data cubes by new variables computed on-the-fly (#272). You can now add a section `Augmentation` into a dataset descriptor, e.g.:
  
```yaml 
Datasets:
  - Identifier: abc
    ...
    Augmentation:
        Path: compute_new_vars.py
        Function: compute_variables
        InputParameters:
            ...
  - ...
```

where `compute_variables` is a function that receives the parent xcube dataset and is expected to return a new dataset with new variables. 

## Note

I've enhanced the demo service configuration in `examples/serve/demo/config.yml` which now computes a new variable (a silly sum of CHL and TSM).



Patch:
diff --git a/CHANGES.md b/CHANGES.md
index 219ffd7e..571eab2d 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -39,6 +39,25 @@
       ...        
   ```      
 
+* The configuration of `xcube serve` has been enhanced to support
+  augmentation of data cubes by new variables computed on-the-fly (#272).
+  You can now add a section `Augmentation` into a dataset descriptor, e.g.:
+  
+  ```yaml 
+    Datasets:
+      - Identifier: abc
+        ...
+        Augmentation:
+          Path: compute_new_vars.py
+          Function: compute_variables
+          InputParameters:
+            ...
+      - ...
+  ```
+  
+  where `compute_variables` is a function that receives the parent xcube dataset
+  and is expected to return a new dataset with new variables. 
+  
 * `xcube serve` now provides basic access control via OAuth2 bearer tokens (#263).
   To configure a service instance with access control, add the following to the 
   `xcube serve` configuration file:
diff --git a/examples/serve/demo/compute_extra_vars.py b/examples/serve/demo/compute_extra_vars.py
new file mode 100644
index 00000000..31a73446
--- /dev/null
+++ b/examples/serve/demo/compute_extra_vars.py
@@ -0,0 +1,9 @@
+import xarray as xr
+
+
+def compute_variables(ds, factor_chl, factor_tsm):
+    chl_tsm_sum = factor_chl * ds.conc_chl + factor_tsm * ds.conc_tsm
+    chl_tsm_sum.attrs.update(dict(units='-',
+                                  long_name='Weighted sum of CHL nd TSM conc.',
+                                  description='Nonsense variable, for demo purpose only'))
+    return xr.Dataset(dict(chl_tsm_sum=chl_tsm_sum))
diff --git a/examples/serve/demo/config.yml b/examples/serve/demo/config.yml
index 53f2110c..3e68f27f 100644
--- a/examples/serve/demo/config.yml
+++ b/examples/serve/demo/config.yml
@@ -11,6 +11,12 @@ Datasets:
     Path: cube-1-250-250.zarr
     Style: default
     TimeSeriesDataset: local_ts
+    Augmentation:
+      Path: "compute_extra_vars.py"
+      Function: "compute_variables"
+      InputParameters:
+        factor_chl: 0.2
+        factor_tsm: 0.7
     PlaceGroups:
       - PlaceGroupRef: inside-cube
       - PlaceGroupRef: outside-cube
diff --git a/xcube/core/mldataset.py b/xcube/core/mldataset.py
index a0a40ca1..255b7269 100644
--- a/xcube/core/mldataset.py
+++ b/xcube/core/mldataset.py
@@ -119,7 +119,7 @@ class LazyMultiLevelDataset(MultiLevelDataset, metaclass=ABCMeta):
         :return: the dataset for the level at *index*.
         """
 
-    def _get_tile_grid_lazily(self):
+    def _get_tile_grid_lazily(self) -> TileGrid:
         """
         Retrieve, i.e. read or compute, the tile grid used by the multi-level dataset.
 
@@ -134,6 +134,40 @@ class LazyMultiLevelDataset(MultiLevelDataset, metaclass=ABCMeta):
                     dataset.close()
 
 
+class CombinedMultiLevelDataset(LazyMultiLevelDataset):
+    """
+    A multi-level dataset that is a combination of other multi-level datasets.
+
+    :param ml_datasets: The multi-level datasets to be combined. At least two must be provided.
+    :param combiner_function: A function used to combine the datasets. It receives a list of
+        datasets (``xarray.Dataset`` instances) and *combiner_params* as keyword arguments.
+        Defaults to function ``xarray.merge()`` with default parameters.
+    :param combiner_params: Parameters to the *combiner_function* passed as keyword arguments.
+    """
+
+    def __init__(self,
+                 ml_datasets: Sequence[MultiLevelDataset],
+                 combiner_function: Callable = None,
+                 combiner_params: Dict[str, Any] = None):
+        super().__init__()
+        if not ml_datasets or len(ml_datasets) < 2:
+            raise ValueError('ml_datasets must have at least two elements')
+        self._ml_datasets = ml_datasets
+        self._combiner_function = combiner_function or xr.merge
+        self._combiner_params = combiner_params or {}
+
+    def _get_tile_grid_lazily(self) -> TileGrid:
+        return self._ml_datasets[0].tile_grid
+
+    def _get_dataset_lazily(self, index: int, **kwargs) -> xr.Dataset:
+        datasets = [ml_dataset.get_dataset(index) for ml_dataset in self._ml_datasets]
+        return self._combiner_function(datasets, **self._combiner_params)
+
+    def close(self):
+        for ml_dataset in self._ml_datasets:
+            ml_dataset.close()
+
+
 class FileStorageMultiLevelDataset(LazyMultiLevelDataset):
     """
     A stored multi-level dataset whose level datasets are lazily read from storage location.
@@ -320,6 +354,8 @@ class ComputedMultiLevelDataset(LazyMultiLevelDataset):
                  input_ml_dataset_getter: Callable[[str], MultiLevelDataset],
                  input_parameters: Dict[str, Any],
                  exception_type=ValueError):
+
+        input_parameters = input_parameters or {}
         super().__init__(kwargs=input_parameters)
 
         try:
@@ -384,7 +420,7 @@ class ComputedMultiLevelDataset(LazyMultiLevelDataset):
         return assert_cube(computed_value, name=self._ds_id)
 
 
-def _get_dataset_tile_grid(dataset: xr.Dataset, num_levels: int = None):
+def _get_dataset_tile_grid(dataset: xr.Dataset, num_levels: int = None) -> TileGrid:
     geo_extent = get_dataset_bounds(dataset)
     inv_y = float(dataset.lat[0]) < float(dataset.lat[-1])
     width, height, tile_width, tile_height = _get_cube_spatial_sizes(dataset)
diff --git a/xcube/webapi/auth.py b/xcube/webapi/auth.py
index 6d58b3f2..2a8c074e 100644
--- a/xcube/webapi/auth.py
+++ b/xcube/webapi/auth.py
@@ -119,11 +119,8 @@ class AuthMixin:
 
         # TODO: read jwks from cache
         response = requests.get(auth_config.well_known_jwks)
-        print(response.content)
         jwks = json.loads(response.content)
 
-        print(json.dumps(jwks))
-
         rsa_key = {}
         for key in jwks["keys"]:
             if key["kid"] == unverified_header["kid"]:
diff --git a/xcube/webapi/context.py b/xcube/webapi/context.py
index cba626b9..b991ee5e 100644
--- a/xcube/webapi/context.py
+++ b/xcube/webapi/context.py
@@ -24,6 +24,7 @@ import logging
 import os
 import os.path
 import threading
+import uuid
 from typing import Any, Dict, List, Optional, Tuple, Callable, Collection, Set
 from typing import Sequence
 
@@ -33,12 +34,12 @@ import pandas as pd
 import s3fs
 import xarray as xr
 import zarr
-
 from xcube.constants import FORMAT_NAME_LEVELS
 from xcube.constants import FORMAT_NAME_NETCDF4
 from xcube.constants import FORMAT_NAME_ZARR
 from xcube.core.dsio import guess_dataset_format
 from xcube.core.mldataset import BaseMultiLevelDataset
+from xcube.core.mldataset import CombinedMultiLevelDataset
 from xcube.core.mldataset import ComputedMultiLevelDataset
 from xcube.core.mldataset import FileStorageMultiLevelDataset
 from xcube.core.mldataset import MultiLevelDataset
@@ -60,6 +61,8 @@ from xcube.webapi.errors import ServiceResourceNotFoundError
 from xcube.webapi.reqparams import RequestParams
 
 COMPUTE_DATASET = 'compute_dataset'
+COMPUTE_VARIABLES = 'compute_variables'
+
 ALL_PLACES = "all"
 
 _LOG = logging.getLogger('xcube')
@@ -215,7 +218,10 @@ class ServiceContext:
 
     def get_variable_for_z(self, ds_id: str, var_name: str, z_index: int) -> xr.DataArray:
         ml_dataset = self.get_ml_dataset(ds_id)
-        dataset = ml_dataset.get_dataset(ml_dataset.num_levels - 1 - z_index)
+        index = ml_dataset.num_levels - 1 - z_index
+        if index < 0 or index >= ml_dataset.num_levels:
+            raise ServiceResourceNotFoundError(f'Variable "{var_name}" has no z-index {z_index} in dataset "{ds_id}"')
+        dataset = ml_dataset.get_dataset(index)
         if var_name not in dataset:
             raise ServiceResourceNotFoundError(f'Variable "{var_name}" not found in dataset "{ds_id}"')
         return dataset[var_name]
@@ -241,9 +247,7 @@ class ServiceContext:
             ds_id = descriptor.get('Identifier')
             file_system = descriptor.get('FileSystem', 'local')
             if file_system == 'local':
-                local_path = descriptor.get('Path')
-                if not os.path.isabs(local_path):
-                    local_path = os.path.join(self.base_dir, local_path)
+                local_path = self.get_descriptor_path(descriptor, f'dataset descriptor {ds_id}')
                 local_path = os.path.normpath(local_path)
                 if os.path.isdir(local_path):
                     s3_bucket_mapping[ds_id] = local_path
@@ -300,15 +304,38 @@ class ServiceContext:
         return ml_dataset, dataset_descriptor
 
     def _open_ml_dataset(self, dataset_descriptor: DatasetDescriptor) -> MultiLevelDataset:
+        ds_id = dataset_descriptor.get('Identifier')
         fs_type = dataset_descriptor.get('FileSystem', 'local')
         if self._ml_dataset_openers and fs_type in self._ml_dataset_openers:
             ml_dataset_opener = self._ml_dataset_openers[fs_type]
         elif fs_type in _DEFAULT_MULTI_LEVEL_DATASET_OPENERS:
             ml_dataset_opener = _DEFAULT_MULTI_LEVEL_DATASET_OPENERS[fs_type]
         else:
-            ds_id = dataset_descriptor.get('Identifier')
             raise ServiceConfigError(f"Invalid fs={fs_type!r} in dataset descriptor {ds_id!r}")
-        return ml_dataset_opener(self, dataset_descriptor)
+        ml_dataset = ml_dataset_opener(self, dataset_descriptor)
+        augmentation = dataset_descriptor.get('Augmentation')
+        if augmentation:
+            path = self.get_descriptor_path(augmentation,
+                                            f"'Augmentation' of dataset descriptor {ds_id}")
+            input_parameters = augmentation.get('InputParameters')
+            callable_name = augmentation.get('Function', COMPUTE_VARIABLES)
+            with measure_time(tag=f"added augmentation from {path}"):
+                aug_id = uuid.uuid4()
+                aug_inp_id = f'augmentation-input-{aug_id}'
+                aug_inp_descriptor = dict(dataset_descriptor)
+                del aug_inp_descriptor['Augmentation']
+                aug_inp_descriptor['Identifier'] = aug_inp_id
+                self._dataset_cache[aug_inp_id] = ml_dataset, aug_inp_descriptor
+                extra_vars = ComputedMultiLevelDataset(f'augmentation-{aug_id}',
+                                                       path,
+                                                       callable_name,
+                                                       [aug_inp_id],
+                                                       self.get_ml_dataset,
+                                                       input_parameters,
+                                                       exception_type=ServiceConfigError)
+                ml_dataset = CombinedMultiLevelDataset([ml_dataset, extra_vars])
+
+        return ml_dataset
 
     def get_legend_label(self, ds_name: str, var_name: str):
         dataset = self.get_dataset(ds_name)
@@ -402,12 +429,7 @@ class ServiceContext:
             place_group = self._place_group_cache[place_group_id]
         else:
             place_group_title = place_group_descriptor.get("Title", place_group_id)
-
-            place_path_wc = place_group_descriptor.get("Path")
-            if not place_path_wc:
-                raise ServiceConfigError("Missing 'Path' entry in a 'PlaceGroups' item")
-            if not os.path.isabs(place_path_wc):
-                place_path_wc = os.path.join(self._base_dir, place_path_wc)
+            place_path_wc = self.get_descriptor_path(place_group_descriptor, f"'PlaceGroups' item")
             source_paths = glob.glob(place_path_wc)
             source_encoding = place_group_descriptor.get("CharacterEncoding", "utf-8")
 
@@ -566,6 +588,18 @@ class ServiceContext:
         # Note: can be optimized by dict/key lookup
         return next((dsd for dsd in dataset_descriptors if dsd['Identifier'] == ds_name), None)
 
+    def get_descriptor_path(self,
+                            descriptor: Dict[str, Any],
+                            descriptor_name: str,
+                            path_entry_name: str = 'Path',
+                            is_url: bool = False) -> str:
+        path = descriptor.get(path_entry_name)
+        if not path:
+            raise ServiceError(f"Missing entry {path_entry_name!r} in {descriptor_name}")
+        if not is_url and not os.path.isabs(path):
+            path = os.path.join(self._base_dir, path)
+        return path
+
 
 def normalize_prefix(prefix: Optional[str]):
     if not prefix:
@@ -588,11 +622,7 @@ def guess_cube_format(path: str) -> str:
 def open_ml_dataset_from_object_storage(ctx: ServiceContext,
                                         dataset_descriptor: DatasetDescriptor) -> MultiLevelDataset:
     ds_id = dataset_descriptor.get('Identifier')
-
-    path = dataset_descriptor.get('Path')
-    if not path:
-        raise ServiceConfigError(f"Missing 'path' entry in dataset descriptor {ds_id}")
-
+    path = ctx.get_descriptor_path(dataset_descriptor, f"dataset descriptor {ds_id}", is_url=True)
     data_format = dataset_descriptor.get('Format', FORMAT_NAME_ZARR)
 
     s3_client_kwargs = {}
@@ -618,14 +648,7 @@ def open_ml_dataset_from_object_storage(ctx: ServiceContext,
 
 def open_ml_dataset_from_local_fs(ctx: ServiceContext, dataset_descriptor: DatasetDescriptor) -> MultiLevelDataset:
     ds_id = dataset_descriptor.get('Identifier')
-
-    path = dataset_descriptor.get('Path')
-    if not path:
-        raise ServiceConfigError(f"Missing 'path' entry in dataset descriptor {ds_id}")
-
-    if not os.path.isabs(path):
-        path = os.path.join(ctx.base_dir, path)
-
+    path = ctx.get_descriptor_path(dataset_descriptor, f"dataset descriptor {ds_id}")
     data_format = dataset_descriptor.get('Format', guess_cube_format(path))
 
     if data_format == FORMAT_NAME_NETCDF4:
@@ -647,14 +670,7 @@ def open_ml_dataset_from_local_fs(ctx: ServiceContext, dataset_descriptor: Datas
 
 def open_ml_dataset_from_python_code(ctx: ServiceContext, dataset_descriptor: DatasetDescriptor) -> MultiLevelDataset:
     ds_id = dataset_descriptor.get('Identifier')
-
-    path = dataset_descriptor.get('Path')
-    if not path:
-        raise ServiceConfigError(f"Missing 'path' entry in dataset descriptor {ds_id}")
-
-    if not os.path.isabs(path):
-        path = os.path.join(ctx.base_dir, path)
-
+    path = ctx.get_descriptor_path(dataset_descriptor, f"dataset descriptor {ds_id}")
     callable_name = dataset_descriptor.get('Function', COMPUTE_DATASET)
     input_dataset_ids = dataset_descriptor.get('InputDatasets', [])
     input_parameters = dataset_descriptor.get('InputParameters', {})


Test patch:
diff --git a/test/core/test_mldataset.py b/test/core/test_mldataset.py
index 079c9110..038d77e0 100644
--- a/test/core/test_mldataset.py
+++ b/test/core/test_mldataset.py
@@ -5,10 +5,48 @@ import numpy as np
 import pandas as pd
 import xarray as xr
 
-from xcube.core.mldataset import BaseMultiLevelDataset, ComputedMultiLevelDataset
+from xcube.core.mldataset import BaseMultiLevelDataset
+from xcube.core.mldataset import CombinedMultiLevelDataset
+from xcube.core.mldataset import ComputedMultiLevelDataset
 from xcube.util.tilegrid import TileGrid
 
 
+class CombinedMultiLevelDatasetTest(unittest.TestCase):
+    def test_it(self):
+        ml_ds_1 = BaseMultiLevelDataset(_get_test_dataset(('noise_1', 'noise_2')))
+        ml_ds_2 = BaseMultiLevelDataset(_get_test_dataset(('noise_3', 'noise_4')))
+        ml_ds_3 = BaseMultiLevelDataset(_get_test_dataset(('noise_5', 'noise_6')))
+
+        ml_ds = CombinedMultiLevelDataset([ml_ds_1, ml_ds_2, ml_ds_3])
+
+        self.assertEqual(3, ml_ds.num_levels)
+        self.assertEqual(TileGrid(3, 2, 1, 180, 180, (-180, -90, 180, 90), inv_y=False),
+                         ml_ds.tile_grid)
+
+        expected_var_names = {'noise_1', 'noise_2',
+                              'noise_3', 'noise_4',
+                              'noise_5', 'noise_6'}
+
+        ds0 = ml_ds.get_dataset(0)
+        self.assertEqual({'time': 14, 'lat': 720, 'lon': 1440, 'bnds': 2}, ds0.dims)
+        self.assertEqual(expected_var_names, set(map(str, ds0.data_vars)))
+        self.assertTrue(all(v.dims == ('time', 'lat', 'lon') for v in ds0.data_vars.values()))
+
+        ds1 = ml_ds.get_dataset(1)
+        self.assertEqual({'time': 14, 'lat': 360, 'lon': 720}, ds1.dims)
+        self.assertEqual(expected_var_names, set(map(str, ds1.data_vars)))
+        self.assertTrue(all(v.dims == ('time', 'lat', 'lon') for v in ds1.data_vars.values()))
+
+        ds2 = ml_ds.get_dataset(2)
+        self.assertEqual({'time': 14, 'lat': 180, 'lon': 360}, ds2.dims)
+        self.assertEqual(expected_var_names, set(map(str, ds2.data_vars)))
+        self.assertTrue(all(v.dims == ('time', 'lat', 'lon') for v in ds2.data_vars.values()))
+
+        self.assertEqual([ds0, ds1, ds2], ml_ds.datasets)
+
+        ml_ds.close()
+
+
 class BaseMultiLevelDatasetTest(unittest.TestCase):
     def test_it(self):
         ds = _get_test_dataset()
@@ -29,6 +67,8 @@ class BaseMultiLevelDatasetTest(unittest.TestCase):
         self.assertIsNot(ds, ds2)
         self.assertEqual({'time': 14, 'lat': 180, 'lon': 360}, ds2.dims)
 
+        self.assertEqual([ds0, ds1, ds2], ml_ds.datasets)
+
         ml_ds.close()
 
 
@@ -63,11 +103,13 @@ class ComputedMultiLevelDatasetTest(unittest.TestCase):
         ds2 = ml_ds2.get_dataset(2)
         self.assertEqual({'time': 3, 'lat': 180, 'lon': 360}, ds2.dims)
 
+        self.assertEqual([ds0, ds1, ds2], ml_ds2.datasets)
+
         ml_ds1.close()
         ml_ds2.close()
 
 
-def _get_test_dataset():
+def _get_test_dataset(var_names=('noise',)):
     w = 1440
     h = 720
     p = 14
@@ -79,7 +121,7 @@ def _get_test_dataset():
     dx = (x2 - x1) / w
     dy = (y2 - y1) / h
 
-    data_vars = dict(noise=(("time", "lat", "lon"), np.random.rand(p, h, w)))
+    data_vars = {var_name: (("time", "lat", "lon"), np.random.rand(p, h, w)) for var_name in var_names}
 
     lat_bnds = np.array(list(zip(np.linspace(y2, y1 + dy, num=h),
                                  np.linspace(y2 - dy, y1, num=h))))
diff --git a/test/webapi/helpers.py b/test/webapi/helpers.py
index 59d00c23..5d8e4ec9 100644
--- a/test/webapi/helpers.py
+++ b/test/webapi/helpers.py
@@ -9,17 +9,10 @@ from xcube.webapi.errors import ServiceBadRequestError
 from xcube.webapi.reqparams import RequestParams
 
 
-def new_test_service_context(ml_dataset_openers: Dict[str, MultiLevelDatasetOpener] = None) -> ServiceContext:
+def new_test_service_context(config_file_name: str = 'config.yml',
+                             ml_dataset_openers: Dict[str, MultiLevelDatasetOpener] = None) -> ServiceContext:
     ctx = ServiceContext(base_dir=get_res_test_dir(), ml_dataset_openers=ml_dataset_openers)
-    config_file = os.path.join(ctx.base_dir, 'config.yml')
-    with open(config_file) as fp:
-        ctx.config = yaml.safe_load(fp)
-    return ctx
-
-
-def new_demo_service_context(ml_dataset_openers: Dict[str, MultiLevelDatasetOpener] = None) -> ServiceContext:
-    ctx = ServiceContext(base_dir=get_res_demo_dir(), ml_dataset_openers=ml_dataset_openers)
-    config_file = os.path.join(ctx.base_dir, 'config.yml')
+    config_file = os.path.join(ctx.base_dir, config_file_name)
     with open(config_file) as fp:
         ctx.config = yaml.safe_load(fp)
     return ctx
diff --git a/test/webapi/res/test/config-aug.yml b/test/webapi/res/test/config-aug.yml
new file mode 100644
index 00000000..4f95adc4
--- /dev/null
+++ b/test/webapi/res/test/config-aug.yml
@@ -0,0 +1,11 @@
+Datasets:
+  - Identifier: demo-aug
+    Title: xcube-server Demonstration L2C Cube with augmentation
+    Path: ../../../../examples/serve/demo/cube-1-250-250.zarr
+    Style: default
+    Augmentation:
+      Path: script.py
+      Function: compute_variables
+      InputParameters:
+        factor_chl: 0.2
+        factor_tsm: 0.7
diff --git a/test/webapi/res/test/config.yml b/test/webapi/res/test/config.yml
index e4eab5a5..df63287d 100644
--- a/test/webapi/res/test/config.yml
+++ b/test/webapi/res/test/config.yml
@@ -8,6 +8,7 @@ Datasets:
     Title: xcube-server Demonstration L3 Cube
     FileSystem: memory
     Path: script.py
+    Function: compute_dataset
     InputDatasets: ["demo"]
     InputParameters:
       period: "1W"
diff --git a/test/webapi/res/test/script.py b/test/webapi/res/test/script.py
index f7704ef4..3395bc0f 100644
--- a/test/webapi/res/test/script.py
+++ b/test/webapi/res/test/script.py
@@ -11,3 +11,11 @@ def compute_dataset(ds, period='1W', incl_stdev=False):
         return ds_merged
     else:
         return ds.resample(time=period).mean(dim='time')
+
+
+def compute_variables(ds, factor_chl, factor_tsm):
+    chl_tsm_sum = factor_chl * ds.conc_chl + factor_tsm * ds.conc_tsm
+    chl_tsm_sum.attrs.update(dict(units='-',
+                                  long_name='Weighted sum of CHL nd TSM conc.',
+                                  description='Nonsense variable, for demo purpose only'))
+    return xr.Dataset(dict(chl_tsm_sum=chl_tsm_sum))
diff --git a/test/webapi/test_context.py b/test/webapi/test_context.py
index 0c651585..946a626b 100644
--- a/test/webapi/test_context.py
+++ b/test/webapi/test_context.py
@@ -4,6 +4,7 @@ import unittest
 import xarray as xr
 
 from test.webapi.helpers import new_test_service_context
+from xcube.core.mldataset import MultiLevelDataset
 from xcube.webapi.errors import ServiceResourceNotFoundError
 
 
@@ -14,6 +15,25 @@ class ServiceContextTest(unittest.TestCase):
         ds = ctx.get_dataset('demo')
         self.assertIsInstance(ds, xr.Dataset)
 
+        ml_ds = ctx.get_ml_dataset('demo')
+        self.assertIsInstance(ml_ds, MultiLevelDataset)
+        self.assertIs(3, ml_ds.num_levels)
+        self.assertIs(ds, ml_ds.get_dataset(0))
+
+        for var_name in ('conc_chl', 'conc_tsm'):
+            for z in range(ml_ds.num_levels):
+                conc_chl_z = ctx.get_variable_for_z('demo', var_name, z)
+                self.assertIsInstance(conc_chl_z, xr.DataArray)
+            with self.assertRaises(ServiceResourceNotFoundError) as cm:
+                ctx.get_variable_for_z('demo', var_name, 3)
+            self.assertEqual(404, cm.exception.status_code)
+            self.assertEqual(f'Variable "{var_name}" has no z-index 3 in dataset "demo"', cm.exception.reason)
+
+        with self.assertRaises(ServiceResourceNotFoundError) as cm:
+            ctx.get_variable_for_z('demo', 'conc_ys', 0)
+        self.assertEqual(404, cm.exception.status_code)
+        self.assertEqual('Variable "conc_ys" not found in dataset "demo"', cm.exception.reason)
+
         with self.assertRaises(ServiceResourceNotFoundError) as cm:
             ctx.get_dataset('demox')
         self.assertEqual(404, cm.exception.status_code)
@@ -24,6 +44,22 @@ class ServiceContextTest(unittest.TestCase):
         self.assertEqual(404, cm.exception.status_code)
         self.assertEqual('Variable "conc_ys" not found in dataset "demo"', cm.exception.reason)
 
+    def test_get_dataset_with_augmentation(self):
+        ctx = new_test_service_context(config_file_name='config-aug.yml')
+
+        ds = ctx.get_dataset('demo-aug')
+        self.assertIsInstance(ds, xr.Dataset)
+
+        ml_ds = ctx.get_ml_dataset('demo-aug')
+        self.assertIsInstance(ml_ds, MultiLevelDataset)
+        self.assertIs(3, ml_ds.num_levels)
+        self.assertIs(ds, ml_ds.get_dataset(0))
+
+        for var_name in ('conc_chl', 'conc_tsm', 'chl_tsm_sum'):
+            for z in range(ml_ds.num_levels):
+                conc_chl_z = ctx.get_variable_for_z('demo-aug', var_name, z)
+                self.assertIsInstance(conc_chl_z, xr.DataArray)
+
     def test_config_and_dataset_cache(self):
         ctx = new_test_service_context()
         self.assertNotIn('demo', ctx.dataset_cache)