From 442490b2e171450488d3dd369588f20e65958887 Mon Sep 17 00:00:00 2001
From: Hauke Schulz <43613877+observingClouds@users.noreply.github.com>
Date: Wed, 30 Oct 2024 11:53:24 +0100
Subject: [PATCH 1/2] move input coordinates to extra subsection

    - introduce `coordinates` subsection
    - prepare for additional variable subsections
---
 README.md                      | 28 +++++++++++++++------------
 example.danra.yaml             | 14 ++++++++------
 mllam_data_prep/ops/loading.py | 35 +++++++++++++++++-----------------
 3 files changed, 42 insertions(+), 35 deletions(-)

diff --git a/README.md b/README.md
index cb220bc..04289a3 100644
--- a/README.md
+++ b/README.md
@@ -149,13 +149,15 @@ inputs:
     dims: [time, x, y, altitude]
     variables:
       u:
-        altitude:
-          values: [100,]
-          units: m
+        coordinates:
+          altitude:
+            values: [100,]
+            units: m
       v:
-        altitude:
-          values: [100, ]
-          units: m
+        coordinates:
+          altitude:
+            values: [100, ]
+            units: m
     dim_mapping:
       time:
         method: rename
@@ -256,13 +258,15 @@ inputs:
     dims: [time, x, y, altitude]
     variables:
       u:
-        altitude:
-          values: [100,]
-          units: m
+        coordinates:
+          altitude:
+            values: [100,]
+            units: m
       v:
-        altitude:
-          values: [100, ]
-          units: m
+        coordinates:
+          altitude:
+            values: [100, ]
+            units: m
     dim_mapping:
       time:
         method: rename
diff --git a/example.danra.yaml b/example.danra.yaml
index 73aa0df..f554d27 100644
--- a/example.danra.yaml
+++ b/example.danra.yaml
@@ -35,13 +35,15 @@ inputs:
     dims: [time, x, y, altitude]
     variables:
       u:
-        altitude:
-          values: [100,]
-          units: m
+        coordinates:
+          altitude:
+            values: [100,]
+            units: m
       v:
-        altitude:
-          values: [100, ]
-          units: m
+        coordinates:
+          altitude:
+            values: [100, ]
+            units: m
     dim_mapping:
       time:
         method: rename
diff --git a/mllam_data_prep/ops/loading.py b/mllam_data_prep/ops/loading.py
index 955fafd..9e78044 100644
--- a/mllam_data_prep/ops/loading.py
+++ b/mllam_data_prep/ops/loading.py
@@ -25,24 +25,25 @@ def load_and_subset_dataset(fp, variables):
     ds_subset = xr.Dataset()
     ds_subset.attrs.update(ds.attrs)
     if isinstance(variables, dict):
-        for var, coords_to_sample in variables.items():
+        for var, var_metadata in variables.items():
             da = ds[var]
-            for coord, sampling in coords_to_sample.items():
-                coord_values = sampling.values
-                try:
-                    da = da.sel(**{coord: coord_values})
-                except KeyError as ex:
-                    raise KeyError(
-                        f"Could not find the all coordinate values `{coord_values}` in "
-                        f"coordinate `{coord}` in the dataset"
-                    ) from ex
-                expected_units = sampling.units
-                coord_units = da[coord].attrs.get("units", None)
-                if coord_units is not None and coord_units != expected_units:
-                    raise ValueError(
-                        f"Expected units {expected_units} for coordinate {coord}"
-                        f" in variable {var} but got {coord_units}"
-                    )
+            if var_metadata.get("coordinates", False):
+                for coord, sampling in var_metadata["coordinates"].items():
+                    coord_values = sampling.values
+                    try:
+                        da = da.sel(**{coord: coord_values})
+                    except KeyError as ex:
+                        raise KeyError(
+                            f"Could not find the all coordinate values `{coord_values}` in "
+                            f"coordinate `{coord}` in the dataset"
+                        ) from ex
+                    expected_units = sampling.units
+                    coord_units = da[coord].attrs.get("units", None)
+                    if coord_units is not None and coord_units != expected_units:
+                        raise ValueError(
+                            f"Expected units {expected_units} for coordinate {coord}"
+                            f" in variable {var} but got {coord_units}"
+                        )
             ds_subset[var] = da
     elif isinstance(variables, list):
         try:

From ff097079a41c95f647b3d891fecba83561c06746 Mon Sep 17 00:00:00 2001
From: Hauke Schulz <43613877+observingClouds@users.noreply.github.com>
Date: Wed, 30 Oct 2024 11:58:59 +0100
Subject: [PATCH 2/2] update schema_version in tests

---
 tests/test_config.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/tests/test_config.py b/tests/test_config.py
index 5459db6..d19c162 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -4,7 +4,7 @@
 import mllam_data_prep as mdp
 
 INVALID_EXTRA_FIELDS_CONFIG_YAML = """
-schema_version: v0.1.0
+schema_version: v0.2.0
 dataset_version: v0.1.0
 
 output:
@@ -21,7 +21,7 @@
 """
 
 MISSING_FIELDS_CONFIG_YAML = """
-schema_version: v0.1.0
+schema_version: v0.2.0
 dataset_version: v0.1.0
 """
 
@@ -36,7 +36,7 @@ def test_get_config_issues():
 
 
 VALID_EXAMPLE_CONFIG_YAML = """
-schema_version: v0.1.0
+schema_version: v0.2.0
 dataset_version: v0.1.0
 
 output:
@@ -71,13 +71,15 @@ def test_get_config_issues():
     dims: [time, x, y, altitude]
     variables:
       u:
-        altitude:
-          values: [100, ]
-          units: m
+        coordinates:
+          altitude:
+            values: [100, ]
+            units: m
       v:
-        altitude:
-          values: [100, ]
-          units: m
+        coordinates:
+          altitude:
+            values: [100, ]
+            units: m
     dim_mapping:
       time:
         method: rename