Fixes for DANRA example (#4)

* minor changes

* put DANRA lsm on S3 test-data bucket

* Add test for DANRA example config file

* run pytest in cicd

* run pytest in cicd

* explicitly install pytest in cicd

---------

Co-authored-by: Leif Denby <[email protected]>

.github/workflows/python-package-pip.yml

@@ -1,5 +1,12 @@
-name: mllam-data-prep (pip install)
-on: [push, pull_request]
+name: pytest
+
+on:
+ push:
+    branches:
+    - master
+ pull_request:
+    branches:
+    - master
 
 jobs:
   test:
@@ -14,4 +21,9 @@ jobs:
       - uses: actions/checkout@v2
       - name: Install package with pip
         run: |
-          pip install .
+          python -m pip install .
+          python -m pip install pytest
+
+      - name: Run tests
+        run: |
+          python -m pytest tests/

example.danra.yaml

@@ -4,13 +4,13 @@ dataset_version: v0.1.0
 architecture:
   sampling_dim: time
   input_variables:
-    static: [grid_index, feature]
+    static: [grid_index, static_feature]
     state: [time, grid_index, state_feature]
     forcing: [time, grid_index, forcing_feature]
   input_range:
     time:
       start: 1990-09-03T00:00
-      end: 1990-09-04T00:00
+      end: 1990-09-09T00:00
       step: PT3H
   chunking:
     time: 6
@@ -43,6 +43,8 @@ inputs:
     path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/single_levels.zarr
     dims: [time, x, y]
     variables:
+      # shouldn't really be using sea-surface pressure as "forcing", but don't
+      # have radiation varibles in danra yet
       - pres_seasurface
     dim_mapping:
       time: time
@@ -53,3 +55,17 @@ inputs:
         method: stack_variables_by_var_name
         name_format: f"{var_name}"
     target_architecture_variable: forcing
+
+  danra_lsm:
+    path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/lsm.zarr
+    dims: [x, y]
+    variables:
+      - lsm
+    dim_mapping:
+      grid_index:
+        method: flatten
+        dims: [x, y]
+      static_feature:
+        method: stack_variables_by_var_name
+        name_format: f"{var_name}"
+    target_architecture_variable: static

mllam_data_prep/create_dataset.py

@@ -139,8 +139,13 @@ def main(fp_config):
 
         da_target.attrs["source_dataset"] = dataset_name
 
+        # only need to do selection for the coordinates that the input dataset actually has
         if architecture_input_ranges is not None:
-            da_target = select_by_kwargs(da_target, **architecture_input_ranges)
+            selection_kwargs = {}
+            for dim in arch_dims:
+                if dim in architecture_input_ranges:
+                    selection_kwargs[dim] = architecture_input_ranges[dim]
+            da_target = select_by_kwargs(da_target, **selection_kwargs)
 
         dataarrays_by_target[target_arch_var].append(da_target)
 

mllam_data_prep/ops/loading.py

@@ -40,7 +40,13 @@ def load_and_subset_dataset(fp, variables):
                     )
             ds_subset[var] = da
     elif isinstance(variables, list):
-        ds_subset = ds[variables]
+        try:
+            ds_subset = ds[variables]
+        except KeyError as ex:
+            raise KeyError(
+                f"Could not find the all variables `{variables}` in the dataset. "
+                f"The available variables are {list(ds.data_vars)}"
+            ) from ex
     else:
         raise ValueError("The `variables` argument should be a list or a dictionary")
     return ds_subset

tests/test_from_config.py

@@ -244,3 +244,8 @@ def test_feature_collision(use_common_feature_var_name):
             mdp.main(fp_config=fp_config)
     else:
         mdp.main(fp_config=fp_config)
+
+
+def test_danra_example():
+    fp_config = Path(__file__).parent.parent / "example.danra.yaml"
+    mdp.main(fp_config=fp_config)