diff --git a/README.md b/README.md index 25e2a90..d797247 100644 --- a/README.md +++ b/README.md @@ -364,7 +364,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 @@ -378,6 +378,97 @@ The `inputs` section defines the source datasets to extract data from. Each sour `start`: start value of the variable `end`: end value of the variable +#### 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). diff --git a/mllam_data_prep/config.py b/mllam_data_prep/config.py index 8a7ccfd..911418a 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,51 @@ 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() + ): + 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{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" + " 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: diff --git a/tests/test_config.py b/tests/test_config.py index 5f7896a..bd0d1b4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,8 +1,12 @@ 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 @@ -129,3 +133,144 @@ def test_config_roundtrip(): 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: + 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 update_config(config: str, update: str): + """ + Update provided config + + Parameters + ---------- + config: str + String with config in yaml format + update: dict + Dictionary with the updated config + + Returns + ------- + config: Config + Updated config + """ + original_config = mdp.Config.from_yaml(config) + update = yaml.safe_load(update) + modified_config = original_config.to_dict() + modified_config.update(update) + + return modified_config + + +@pytest.mark.parametrize( + "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 + ), + ( + # 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`", + ), + ( + # 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", + ), + ], +) +def test_config_validation_of_inputs_section( + base_config, 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 + """ + invalid_config_dict = update_config(base_config, invalid_inputs_section) + + with pytest.raises(expected_exception, match=expected_message): + mdp.Config.from_dict(invalid_config_dict) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index e78a93b..7770a62 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. @@ -255,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", [