From 86be3c40ec7f17db61497b5935d2839c7cd3eb38 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 7 Feb 2025 10:15:56 +0000 Subject: [PATCH 1/8] Add check in the config validation for the selection of variables from different coord levels in the same input dataset, since this is not yet supported --- mllam_data_prep/config.py | 52 ++++++++++++++++++++++++++++++++++++--- 1 file changed, 48 insertions(+), 4 deletions(-) diff --git a/mllam_data_prep/config.py b/mllam_data_prep/config.py index dbf165c..e0a47c6 100644 --- a/mllam_data_prep/config.py +++ b/mllam_data_prep/config.py @@ -11,10 +11,14 @@ class InvalidConfigException(Exception): def validate_config(config_inputs): """ - Validate that, in the config: - - either `variables` or `derived_variables` are present in the config - - if both `variables` and `derived_variables` are present, that they don't - add the same variables to the dataset + Validate that, in a dataset in the inputs section of the config: + 1. at least either `variables` or `derived_variables` are present in the config + 2. if both `variables` and `derived_variables` are present, that they don't + include overlapping variable names + 3. if the `variables` section contains a selection of variables by coordinates + that all variables are selected on the same coordinate since we currently + do not support variables selected from different coord levels in the same inputs + dataset Parameters ---------- @@ -25,13 +29,17 @@ def validate_config(config_inputs): """ for input_dataset_name, input_dataset in config_inputs.items(): + # Check that at least either `variables` or `derived_variables` are preset if not input_dataset.variables and not input_dataset.derived_variables: raise InvalidConfigException( + "Missing `variables` and/or `derived_variables`\n." f"Input dataset '{input_dataset_name}' is missing the keys `variables` and/or" " `derived_variables`. Make sure that you update the config so that the input" f" dataset '{input_dataset_name}' contains at least either a `variables` or" " `derived_variables` section." ) + # Check that there are no overlapping variable names between `variables` + # and `derived_variables` elif input_dataset.variables and input_dataset.derived_variables: # Check so that there are no overlapping variables if isinstance(input_dataset.variables, list): @@ -46,12 +54,48 @@ def validate_config(config_inputs): common_vars = list(set(variable_vars) & set(derived_variable_vars)) if len(common_vars) > 0: raise InvalidConfigException( + "Overlapping variable names in `variables` and `derived_variables`\n." "Both `variables` and `derived_variables` include the following variables name(s):" f" '{', '.join(common_vars)}'. This is not allowed. Make sure that there" " are no overlapping variable names between `variables` and `derived_variables`," f" either by renaming or removing '{', '.join(common_vars)}' from one of them." ) + # Check the values of the selected coordinates among variables + # We are currently not supporting the selection of multiple variables from different + # coord levels. + dim_mapping = input_dataset.dim_mapping.copy() + variable_dim_mappings = {} + for arch_dim in list(dim_mapping.keys()): + if dim_mapping[arch_dim].method == "stack_variables_by_var_name": + variable_dim_mappings[arch_dim] = dim_mapping.pop(arch_dim) + + if input_dataset.variables and isinstance(input_dataset.variables, dict): + dict_of_coord_values = {} + for arch_dim, input_dim_map in variable_dim_mappings.items(): + dims = input_dim_map.dims[0] + + for var_name, var in input_dataset.variables.items(): + dict_of_coord_values[var_name] = var[dims].values + + if len(input_dataset.variables) > 1: + if not all( + coord_values == list(dict_of_coord_values.values())[0] + for coord_values in dict_of_coord_values.values() + ): + raise NotImplementedError( + "Selection of variables from different coord levels is currently not supported.\n" + f"Input dataset '{input_dataset_name}' is trying to select variables" + f" '{', '.join(list(input_dataset.variables.keys()))}' from different coord levels.\n" + f"Coord values per variable to select from:\n" + f"{'\n'.join(f' - {var}: {values}' for var, values in dict_of_coord_values.items())}" + "\nThis type of selection of variables with different coord levels is currently not supported." + "\nIf you want to make this type of selection you can split your selection among several" + " input datasets. Either you split it such that each input dataset is only containing a" + " selection of variables from the same coord levels, or such that each input dataset is only" + " containing one variable with the selection from the desired coord levels." + ) + @dataclass class Range: From 45678bde2a00a5fdea091496678509d87bfbfa06 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 7 Feb 2025 10:23:52 +0000 Subject: [PATCH 2/8] Add tests for checking that the exceptions in the config validation are caught --- tests/test_config.py | 159 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 159 insertions(+) diff --git a/tests/test_config.py b/tests/test_config.py index 5459db6..56bca1e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,3 +1,8 @@ +import re +import shutil +import tempfile +from pathlib import Path + import pytest from dataclass_wizard.errors import MissingFields, UnknownJSONKey @@ -119,3 +124,157 @@ def test_get_config_nested(): assert input_config.target_output_variable is not None with pytest.raises(AttributeError): input_config.foobarfield + + +INVALID_MISSING_VARIABLES_YAML = """ +inputs: + danra_height_levels: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/single_levels.zarr + dims: [time, x, y] + dim_mapping: + time: + method: rename + dim: time + state_feature: + method: stack_variables_by_var_name + name_format: "{var_name}" + grid_index: + method: stack + dims: [x, y] + target_output_variable: state +""" + +INVALID_OVERLAPPING_VARIABLE_NAMES_YAML = """ +inputs: + danra_height_levels: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/single_levels.zarr + dims: [time, x, y] + variables: + - swavr0m + derived_variables: + swavr0m: + kwargs: + time: ds_input.time + lat: ds_input.lat + lon: ds_input.lon + function: mllam_data_prep.ops.derive_variable.physical_field.calculate_toa_radiation + dim_mapping: + time: + method: rename + dim: time + state_feature: + method: stack_variables_by_var_name + name_format: "{var_name}" + grid_index: + method: stack + dims: [x, y] + target_output_variable: state +""" + +INVALID_INPUTS_COORDS_SELECTION_YAML = """ +inputs: + danra_height_levels: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/height_levels.zarr + dims: [time, x, y, altitude] + variables: + u: + altitude: + values: [100, ] + units: m + v: + altitude: + values: [50, ] + units: m + dim_mapping: + time: + method: rename + dim: time + state_feature: + method: stack_variables_by_var_name + dims: [altitude] + name_format: "{var_name}{altitude}m" + grid_index: + method: stack + dims: [x, y] + target_output_variable: state +""" + + +def modify_example_config_inputs_section_yaml(new_inputs_section): + """ + Get the example config file as a yaml string and replace the + `inputs` section with a new inputs section before + + Parameters + ---------- + new_inputs_section: str + String with a new inputs section + Returns + ------- + modified_yaml: str + Modified config yaml with the new inputs section replacing + the old one in the example config + """ + # Copy the config file to a temporary directory before reading it + fp_example = "example.danra.yaml" + tmpdir = tempfile.TemporaryDirectory() + fp_config_copy = Path(tmpdir.name) / fp_example + shutil.copy(fp_example, fp_config_copy) + + # Read the example config file as text to preserve the order + base_config_yaml = Path(fp_config_copy).read_text() + + # Use regex to replace the entire "inputs:" block + if new_inputs_section: + modified_yaml = re.sub( + r"inputs:\n((?:\s{2,}.*\n)*)", # Matches "inputs:" and all indented content + new_inputs_section, + base_config_yaml, + ) + else: + modified_yaml = base_config_yaml + + return modified_yaml + + +@pytest.mark.parametrize( + "invalid_inputs_section, expected_exception, expected_message", + [ + ( + # Invalid inputs section with missing `variables` and/or `derived_variables` + INVALID_MISSING_VARIABLES_YAML, # invalid config yaml example + mdp.InvalidConfigException, # expected exception + "Missing `variables` and/or `derived_variables`", # expected error message from + ), + ( + # Invalid inputs section with overlapping variable names between `variables` + # and `derived_variables` + INVALID_OVERLAPPING_VARIABLE_NAMES_YAML, + mdp.InvalidConfigException, + "Overlapping variable names in `variables` and `derived_variables`", + ), + ( + # Invalid inputs section selection of variables from different coords levels, + # which is currently not implemented + INVALID_INPUTS_COORDS_SELECTION_YAML, + NotImplementedError, + "Selection of variables from different coord levels is currently not supported", + ), + ], +) +def test_config_validation_of_inputs_section( + invalid_inputs_section, expected_exception, expected_message +): + """ + Test that `validate_config` raises the coerrect exceptions when an inputs dataset + 1. is missing `variables` and/or `derived_variables` + 2. has overlapping variable names between `variables` and `derived_variables` + 3. includes a selection of variables from different coord levels, since this is not yet implemented + """ + # Modify the example config + invalid_config_yaml = modify_example_config_inputs_section_yaml( + invalid_inputs_section + ) + + with pytest.raises(expected_exception, match=expected_message): + mdp.Config.from_yaml(invalid_config_yaml) From 7cc73b7bdf25c35f026336d608bfe01422b5c1e8 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 7 Feb 2025 12:18:45 +0000 Subject: [PATCH 3/8] Add section about retrictions on selection on coordinates to README --- README.md | 93 ++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 92 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 0899771..560cf56 100644 --- a/README.md +++ b/README.md @@ -350,7 +350,7 @@ The `inputs` section defines the source datasets to extract data from. Each sour - `path`: the path to the source dataset. This can be a local path or a URL to e.g. a zarr dataset or netCDF file, anything that can be read by `xarray.open_dataset(...)`. - `dims`: the dimensions that the source dataset is expected to have. This is used to check that the source dataset has the expected dimensions and also makes it clearer in the config file what the dimensions of the source dataset are. -- `variables`: selects which variables to extract from the source dataset. This may either be a list of variable names, or a dictionary where each key is the variable name and the value defines a dictionary of coordinates to do selection on. When doing selection you may also optionally define the units of the variable to check that the units of the variable match the units of the variable in the model architecture. +- `variables`: selects which variables to extract from the source dataset. This may either be a list of variable names, or a dictionary where each key is the variable name and the value defines a dictionary of coordinates to do selection on. When doing selection you may also optionally define the units of the variable to check that the units of the variable match the units of the variable in the model architecture. See also the 'Selection of variables by coordinates' section for more details on the type of selections allowed. - `target_output_variable`: the variable in the model architecture that the source dataset should be mapped to. - `dim_mapping`: defines how the dimensions of the source dataset should be mapped to the dimensions of the model architecture. This is done by defining a method to apply to each dimension. The methods are: - `rename`: simply rename the dimension to the new name @@ -360,6 +360,97 @@ The `inputs` section defines the source datasets to extract data from. Each sour - `function`: the function used to derive a variable. This should be a string with the full namespace of the function, e.g. `mllam_data_prep.ops.derived_variables.physical_field.calculate_toa_radiation`. - `kwargs`: arguments to `function`. This is a dictionary where each key is the named argument to `function` and each value is the input to the function. Here we distinguish between values to be extracted/selected from the input dataset and values supplied by the users themselves. Arguments with values to be extracted from the input dataset need to be prefixed with "ds_input." to distinguish them from other arguments. See the 'Derived Variables' section for more details. +#### Selection of variables by coordinates +When selecting variables to extract from the source dataset it is possible to also specify a selection on coordinates. Currently it is only allowed to, within one input dataset, specify the variables to extract with a selection on coordinates **if the coordinates are for the same levels**, as illustrated in the example config [example.danra.yaml](example.danra.yaml), and also reproduced below, for the selection of `u` and `v` at 100 m height. +```yaml + variables: + u: + altitude: + values: [100,] + units: m + v: + altitude: + values: [100, ] + units: m +``` + +Currently, support for the selection of variables from different coordinate levels within one input dataset is not implemented. This means that the following selections are currently not allowed within the same input dataset +```yaml + variables: + z: + altitude: + values: [100, 50] + units: m + t: + altitude: + values: [100, ] + units: m +``` +and +```yaml + variables: + z: + altitude: + values: [50] + units: m + t: + altitude: + values: [100, ] + units: m +``` + +Instead you need to split the selection across multiple input datasets. For the first example, you need have two separate input datasets, where you could split it in two separate ways. You could either do +```yaml +inputs: + danra_height_levels_100m: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/height_levels.zarr + dims: [time, x, y, altitude] + variables: + z: + altitude: + values: [100,] + units: m + t: + altitude: + values: [100, ] + units: m + ... + + danra_height_levels_50m: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/height_levels.zarr + dims: [time, x, y, altitude] + variables: + z: + altitude: + values: [50,] + units: m + ... +``` +or +```yaml +inputs: + danra_height_levels_z: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/height_levels.zarr + dims: [time, x, y, altitude] + variables: + z: + altitude: + values: [100, 50,] + units: m + ... + + danra_height_levels_t: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/height_levels.zarr + dims: [time, x, y, altitude] + variables: + t: + altitude: + values: [100,] + units: m + ... +``` +Similarly, you need to split the selection in the second example across two separate input datasets. + #### Derived Variables Variables that are not part of the source dataset but can be derived from variables in the source dataset can also be included. They should be defined in their own section, called `derived_variables` as illustrated in the example config above and in the example config file [example.danra.yaml](example.danra.yaml). From 2667d842ee68567ed8b926d68e5ea5b6b7fdd161 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 7 Feb 2025 14:34:48 +0100 Subject: [PATCH 4/8] Fix error message --- mllam_data_prep/config.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/mllam_data_prep/config.py b/mllam_data_prep/config.py index e0a47c6..47d61a5 100644 --- a/mllam_data_prep/config.py +++ b/mllam_data_prep/config.py @@ -83,12 +83,15 @@ def validate_config(config_inputs): coord_values == list(dict_of_coord_values.values())[0] for coord_values in dict_of_coord_values.values() ): + formatted_variables_and_coord_selection = "\n".join( + f" - {var}: {values}" + for var, values in dict_of_coord_values.items() + ) raise NotImplementedError( "Selection of variables from different coord levels is currently not supported.\n" f"Input dataset '{input_dataset_name}' is trying to select variables" f" '{', '.join(list(input_dataset.variables.keys()))}' from different coord levels.\n" - f"Coord values per variable to select from:\n" - f"{'\n'.join(f' - {var}: {values}' for var, values in dict_of_coord_values.items())}" + f"Coord values per variable to select from:\n{formatted_variables_and_coord_selection}" "\nThis type of selection of variables with different coord levels is currently not supported." "\nIf you want to make this type of selection you can split your selection among several" " input datasets. Either you split it such that each input dataset is only containing a" From 1c1dc4ad8db13e8bd0c183a7b4386f932475faa4 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 14 Feb 2025 13:02:45 +0000 Subject: [PATCH 5/8] Update function modifying the base config --- tests/test_config.py | 64 +++++++++++++++++--------------------------- 1 file changed, 25 insertions(+), 39 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 62032b7..a43499b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,13 +1,12 @@ -import re -import shutil -import tempfile -from pathlib import Path - import pytest +import yaml from dataclass_wizard.errors import MissingFields, UnknownJSONKey import mllam_data_prep as mdp +with open("example.danra.yaml", "r") as file: + BASE_CONFIG = file.read() + INVALID_EXTRA_FIELDS_CONFIG_YAML = """ schema_version: v0.1.0 dataset_version: v0.1.0 @@ -200,48 +199,36 @@ def test_get_config_nested(): """ -def modify_example_config_inputs_section_yaml(new_inputs_section): +def update_config(config: str, update: str): """ - Get the example config file as a yaml string and replace the - `inputs` section with a new inputs section before + Update provided config Parameters ---------- - new_inputs_section: str - String with a new inputs section + config: str + String with config in yaml format + update: dict + Dictionary with the updated config + Returns ------- - modified_yaml: str - Modified config yaml with the new inputs section replacing - the old one in the example config + config: Config + Updated config """ - # Copy the config file to a temporary directory before reading it - fp_example = "example.danra.yaml" - tmpdir = tempfile.TemporaryDirectory() - fp_config_copy = Path(tmpdir.name) / fp_example - shutil.copy(fp_example, fp_config_copy) - - # Read the example config file as text to preserve the order - base_config_yaml = Path(fp_config_copy).read_text() - - # Use regex to replace the entire "inputs:" block - if new_inputs_section: - modified_yaml = re.sub( - r"inputs:\n((?:\s{2,}.*\n)*)", # Matches "inputs:" and all indented content - new_inputs_section, - base_config_yaml, - ) - else: - modified_yaml = base_config_yaml + original_config = mdp.Config.from_yaml(config) + update = yaml.safe_load(update) + modified_config = original_config.to_dict() + modified_config.update(update) - return modified_yaml + return modified_config @pytest.mark.parametrize( - "invalid_inputs_section, expected_exception, expected_message", + "base_config, invalid_inputs_section, expected_exception, expected_message", [ ( # Invalid inputs section with missing `variables` and/or `derived_variables` + BASE_CONFIG, INVALID_MISSING_VARIABLES_YAML, # invalid config yaml example mdp.InvalidConfigException, # expected exception "Missing `variables` and/or `derived_variables`", # expected error message from @@ -249,6 +236,7 @@ def modify_example_config_inputs_section_yaml(new_inputs_section): ( # Invalid inputs section with overlapping variable names between `variables` # and `derived_variables` + BASE_CONFIG, INVALID_OVERLAPPING_VARIABLE_NAMES_YAML, mdp.InvalidConfigException, "Overlapping variable names in `variables` and `derived_variables`", @@ -256,6 +244,7 @@ def modify_example_config_inputs_section_yaml(new_inputs_section): ( # Invalid inputs section selection of variables from different coords levels, # which is currently not implemented + BASE_CONFIG, INVALID_INPUTS_COORDS_SELECTION_YAML, NotImplementedError, "Selection of variables from different coord levels is currently not supported", @@ -263,7 +252,7 @@ def modify_example_config_inputs_section_yaml(new_inputs_section): ], ) def test_config_validation_of_inputs_section( - invalid_inputs_section, expected_exception, expected_message + base_config, invalid_inputs_section, expected_exception, expected_message ): """ Test that `validate_config` raises the coerrect exceptions when an inputs dataset @@ -271,13 +260,10 @@ def test_config_validation_of_inputs_section( 2. has overlapping variable names between `variables` and `derived_variables` 3. includes a selection of variables from different coord levels, since this is not yet implemented """ - # Modify the example config - invalid_config_yaml = modify_example_config_inputs_section_yaml( - invalid_inputs_section - ) + invalid_config_dict = update_config(base_config, invalid_inputs_section) with pytest.raises(expected_exception, match=expected_message): - mdp.Config.from_yaml(invalid_config_yaml) + mdp.Config.from_dict(invalid_config_dict) def test_config_roundtrip(): From 6a67625cdfb082a5ea90e476a90ceb68dc5f9cd8 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 14 Feb 2025 13:04:07 +0000 Subject: [PATCH 6/8] Add skip on NotImplementedError functionality --- tests/test_dataset.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index e78a93b..2e375e9 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,4 +1,6 @@ """Tests for the output dataset created by `mllam-data-prep`.""" +from functools import wraps + import pytest import yaml @@ -139,6 +141,25 @@ """ +def skip_on(exception, reason="Default reason"): + """Wrapper for skipping tests based on specified exception.""" + # Func below is the real decorator and will receive the test function as parameter + def decorator_func(f): + @wraps(f) + def wrapper(*args, **kwargs): + try: + # Try to run the test + return f(*args, **kwargs) + except exception: + # If exception of given type happens + # just swallow it and raise pytest.Skip with given reason + pytest.skip(reason) + + return wrapper + + return decorator_func + + def update_config(config: str, update: str): """ Update provided config. @@ -164,6 +185,7 @@ def update_config(config: str, update: str): return modified_config +@skip_on(NotImplementedError, reason="This functionality is not yet implemented.") @pytest.mark.parametrize( "base_config, new_inputs_section", [ From 46036bd13ef196b709050704a878e1ba415d62de Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 14 Feb 2025 13:05:57 +0000 Subject: [PATCH 7/8] Restructure file --- tests/test_config.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index a43499b..bd0d1b4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -125,6 +125,16 @@ def test_get_config_nested(): input_config.foobarfield +def test_config_roundtrip(): + original_config = mdp.Config.from_yaml(VALID_EXAMPLE_CONFIG_YAML) + roundtrip_config_dict = mdp.Config.from_dict(original_config.to_dict()) + roundtrip_config_yaml = mdp.Config.from_yaml(original_config.to_yaml()) + roundtrip_config_json = mdp.Config.from_json(original_config.to_json()) + assert original_config == roundtrip_config_dict + assert original_config == roundtrip_config_yaml + assert original_config == roundtrip_config_json + + INVALID_MISSING_VARIABLES_YAML = """ inputs: danra_height_levels: @@ -264,13 +274,3 @@ def test_config_validation_of_inputs_section( with pytest.raises(expected_exception, match=expected_message): mdp.Config.from_dict(invalid_config_dict) - - -def test_config_roundtrip(): - original_config = mdp.Config.from_yaml(VALID_EXAMPLE_CONFIG_YAML) - roundtrip_config_dict = mdp.Config.from_dict(original_config.to_dict()) - roundtrip_config_yaml = mdp.Config.from_yaml(original_config.to_yaml()) - roundtrip_config_json = mdp.Config.from_json(original_config.to_json()) - assert original_config == roundtrip_config_dict - assert original_config == roundtrip_config_yaml - assert original_config == roundtrip_config_json From e47600643f0a30949552896d5d94c5051dbd1f98 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 14 Feb 2025 13:10:03 +0000 Subject: [PATCH 8/8] Put the 'skip_on' decorator on the correct test --- tests/test_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 2e375e9..7770a62 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -185,7 +185,6 @@ def update_config(config: str, update: str): return modified_config -@skip_on(NotImplementedError, reason="This functionality is not yet implemented.") @pytest.mark.parametrize( "base_config, new_inputs_section", [ @@ -277,6 +276,7 @@ def test_selected_output_variables(base_config, new_inputs_section): pytest.fail(error_message) +@skip_on(NotImplementedError, reason="This functionality is not yet implemented.") @pytest.mark.parametrize( "base_config, update, expected_result", [