From 7a7fc0ea6a8eb7371f6a8e09541d4d119aace892 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Tue, 4 Feb 2025 07:03:14 +0000 Subject: [PATCH 1/5] Apply fix 8040378 from joeloskarsson --- mllam_data_prep/create_dataset.py | 5 ++--- mllam_data_prep/ops/derive_variable/main.py | 7 ++++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mllam_data_prep/create_dataset.py b/mllam_data_prep/create_dataset.py index b5d06df..3618e77 100644 --- a/mllam_data_prep/create_dataset.py +++ b/mllam_data_prep/create_dataset.py @@ -154,11 +154,9 @@ def create_dataset(config: Config): except Exception as ex: raise Exception(f"Error loading dataset {dataset_name} from {path}") from ex - # Initialize the output dataset and add dimensions + # Initialize the output dataset ds = xr.Dataset() ds.attrs.update(ds_input.attrs) - for dim in ds_input.dims: - ds = ds.assign_coords({dim: ds_input.coords[dim]}) if selected_variables: logger.info(f"Extracting selected variables from dataset {dataset_name}") @@ -184,6 +182,7 @@ def create_dataset(config: Config): ds=ds_input, derived_variable=derived_variable, chunking=chunking_config, + target_dims=expected_input_var_dims, ) _check_dataset_attributes( diff --git a/mllam_data_prep/ops/derive_variable/main.py b/mllam_data_prep/ops/derive_variable/main.py index 22a3186..6f42f1c 100644 --- a/mllam_data_prep/ops/derive_variable/main.py +++ b/mllam_data_prep/ops/derive_variable/main.py @@ -18,7 +18,7 @@ REQUIRED_FIELD_ATTRIBUTES = ["units", "long_name"] -def derive_variable(ds, derived_variable, chunking): +def derive_variable(ds, derived_variable, chunking, target_dims): """ Derive a variable using the `function` and `kwargs` of `derived_variable`. @@ -33,6 +33,9 @@ def derive_variable(ds, derived_variable, chunking): chunking: Dict[str, int] Dictionary with keys as the dimensions to chunk along and values with the chunk size + target_dims: List[str] + List of dims from ds to broadcast derived variable to, + if not used in calculation Returns ------- @@ -40,8 +43,6 @@ def derive_variable(ds, derived_variable, chunking): Dataset with derived variables included """ - target_dims = list(ds.sizes.keys()) - function_namespace = derived_variable.function expected_field_attributes = derived_variable.attrs From 82f8e59f604a5218e4fc24cb05c2f6914c374738 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Thu, 6 Feb 2025 13:11:08 +0000 Subject: [PATCH 2/5] Add tests - Test that the output variables match the ones specified in the config - Test that the output dataset does not contain any nans --- tests/test_dataset.py | 300 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 300 insertions(+) create mode 100644 tests/test_dataset.py diff --git a/tests/test_dataset.py b/tests/test_dataset.py new file mode 100644 index 0000000..4033dc9 --- /dev/null +++ b/tests/test_dataset.py @@ -0,0 +1,300 @@ +import re +import shutil +import tempfile +from pathlib import Path + +import pytest + +import mllam_data_prep as mdp + +HEIGHT_LEVEL_TEST_SECTION = """\ +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, 50,] + units: m + v: + altitude: + values: [100, 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 +""" + +PRESSURE_LEVEL_TEST_SECTION = """\ +inputs: + danra_pressure_levels: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/pressure_levels.zarr + dims: [time, x, y, pressure] + variables: + u: + pressure: + values: [1000,] + units: hPa + v: + pressure: + values: [1000, ] + units: hPa + dim_mapping: + time: + method: rename + dim: time + state_feature: + method: stack_variables_by_var_name + dims: [pressure] + name_format: "{var_name}{pressure}m" + grid_index: + method: stack + dims: [x, y] + target_output_variable: state +""" + +SINGLE_LEVEL_SELECTED_VARIABLES_TEST_SECTION = """\ +inputs: + danra_single_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: + - t2m + - pres_seasurface + 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 +""" + +SINGLE_LEVEL_DERIVED_VARIABLES_TEST_SECTION = """\ +inputs: + danra_single_levels: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/single_levels.zarr + dims: [time, x, y] + derived_variables: + # derive variables to be used as forcings + toa_radiation: + 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 +""" + + +def modify_example_config(new_inputs_section): + """ + Get the example config file as a yaml string and replace the + `inputs` section with a new inputs section before reading the config + + Parameters + ---------- + new_inputs_section: str + String with a new inputs section + + Returns + ------- + config: Config + Modified config 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 + + # Read the config + modified_config = mdp.Config.from_yaml(modified_yaml) + + return modified_config + + +@pytest.mark.parametrize( + "new_inputs_section", + [ + None, # Does not modify the example config + PRESSURE_LEVEL_TEST_SECTION, + HEIGHT_LEVEL_TEST_SECTION, + SINGLE_LEVEL_SELECTED_VARIABLES_TEST_SECTION, + SINGLE_LEVEL_DERIVED_VARIABLES_TEST_SECTION, + ], +) +def test_selected_output_variables(new_inputs_section): + """ + Test that the variables specified in each input dataset are + present in the output dataset. + """ + # Modify the example config + config = modify_example_config(new_inputs_section) + + # Create the dataset + ds = mdp.create_dataset(config=config) + + # Check that the output variables are the ones selected + for dataset_name, input_config in config.inputs.items(): + target_output_variable = input_config.target_output_variable + + # Get the expected selected variable names + selected_variables = input_config.variables or [] + if isinstance(selected_variables, dict): + selected_var_names = list(selected_variables.keys()) + elif isinstance(selected_variables, list): + selected_var_names = selected_variables + else: + pytest.fail( + "Expected either 'list' or 'dict' but got" + f" type {type(selected_variables)} for 'selected_variables'." + ) + + # Get the expected derived variable names + derived_variables = input_config.derived_variables or [] + if isinstance(derived_variables, dict): + derived_var_names = list(derived_variables.keys()) + elif isinstance(derived_variables, list): + derived_var_names = derived_variables + else: + pytest.fail( + "Expected either 'list' or 'dict' but got" + f" type {type(derived_variables)} for 'derived_variables'." + ) + + dim_mapping = input_config.dim_mapping[target_output_variable + "_feature"] + dims = dim_mapping.dims or [] + name_format = dim_mapping.name_format + + if len(dims) == 0: + selected_vars = selected_var_names + derived_vars = derived_var_names + elif len(dims) == 1: + coord = dims[0] + # Stack the variable names by coordinates, as is done in + # mdp.ops.stacking.stack_variables_by_coord_values + selected_vars = [] + for var_name in selected_var_names: + coord_values = selected_variables[var_name][coord].values + formatted_var_names = [ + name_format.format(var_name=var_name, **{coord: val}) + for val in coord_values + ] + selected_vars += formatted_var_names + # We currently do not support stacking of variables by coordinates + # for the derived variables + derived_vars = [] + + expected_variables = selected_vars + derived_vars + output_variables = ds[target_output_variable + "_feature"].values + + if set(expected_variables) != set(output_variables): + # Check if there are missing or extra variable + missing_vars = list(set(expected_variables) - set(output_variables)) + extra_vars = list(set(output_variables) - set(expected_variables)) + + error_message = ( + f"Expected {expected_variables}, but got {output_variables}." + ) + if missing_vars: + error_message += f"\nMissing variables: {missing_vars}" + if extra_vars: + error_message += f"\nExtra variables: {extra_vars}" + + pytest.fail(error_message) + + +INVALID_PRESSURE_LEVEL_TEST_SECTION = """\ +inputs: + danra_pressure_levels: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/pressure_levels.zarr + dims: [time, x, y, pressure] + variables: + z: + pressure: + values: [1000,] + units: hPa + t: + pressure: + values: [800, ] + units: hPa + dim_mapping: + time: + method: rename + dim: time + state_feature: + method: stack_variables_by_var_name + dims: [pressure] + name_format: "{var_name}{pressure}m" + grid_index: + method: stack + dims: [x, y] + target_output_variable: state +""" + + +@pytest.mark.parametrize( + "new_inputs_section, expected_result", + [ + ( + None, + False, + ), # Do not modify the example config - should return False since we're expecting no nans + ( + INVALID_PRESSURE_LEVEL_TEST_SECTION, + True, + ), # Dataset with nans - should return True + ], +) +def test_output_dataset_for_nans(new_inputs_section, expected_result): + """ + Test that the output dataset does not contain any nan values. + """ + # Modify the example config + config = modify_example_config(new_inputs_section) + + # Create the dataset + ds = mdp.create_dataset(config=config) + + # Test that we have no nans + nan_in_ds = ds.isnull().any().compute().to_array().any().item() + assert nan_in_ds == expected_result From ecd3e6ecb61d8e0faa748b104ba24ba8db00fb6a Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 7 Feb 2025 07:32:03 +0100 Subject: [PATCH 3/5] Better name for function modifying the inputs section in the example config --- tests/test_dataset.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 4033dc9..ecdcfae 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,3 +1,4 @@ +"""Tests for the output dataset created by `mllam-data-prep`.""" import re import shutil import tempfile @@ -111,7 +112,7 @@ """ -def modify_example_config(new_inputs_section): +def modify_example_config_inputs_section(new_inputs_section): """ Get the example config file as a yaml string and replace the `inputs` section with a new inputs section before reading the config @@ -168,7 +169,7 @@ def test_selected_output_variables(new_inputs_section): present in the output dataset. """ # Modify the example config - config = modify_example_config(new_inputs_section) + config = modify_example_config_inputs_section(new_inputs_section) # Create the dataset ds = mdp.create_dataset(config=config) @@ -186,7 +187,7 @@ def test_selected_output_variables(new_inputs_section): else: pytest.fail( "Expected either 'list' or 'dict' but got" - f" type {type(selected_variables)} for 'selected_variables'." + f" type {type(selected_variables)} for 'variables'." ) # Get the expected derived variable names @@ -290,7 +291,7 @@ def test_output_dataset_for_nans(new_inputs_section, expected_result): Test that the output dataset does not contain any nan values. """ # Modify the example config - config = modify_example_config(new_inputs_section) + config = modify_example_config_inputs_section(new_inputs_section) # Create the dataset ds = mdp.create_dataset(config=config) From e379bb8c2e7baf60f9e23aec67e95db6ba6dcfe0 Mon Sep 17 00:00:00 2001 From: Hauke Schulz <43613877+observingClouds@users.noreply.github.com> Date: Fri, 14 Feb 2025 08:31:06 +0100 Subject: [PATCH 4/5] Pr/60 (#2) Refactor the config modifier to make it more general and make the file more structured --- tests/test_dataset.py | 141 ++++++++++++++++++------------------------ 1 file changed, 60 insertions(+), 81 deletions(-) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index ecdcfae..e78a93b 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,13 +1,12 @@ """Tests for the output dataset created by `mllam-data-prep`.""" -import re -import shutil -import tempfile -from pathlib import Path - import pytest +import yaml import mllam_data_prep as mdp +with open("example.danra.yaml", "r") as file: + BASE_CONFIG = file.read() + HEIGHT_LEVEL_TEST_SECTION = """\ inputs: danra_height_levels: @@ -111,71 +110,83 @@ target_output_variable: state """ +INVALID_PRESSURE_LEVEL_TEST_SECTION = """\ +inputs: + danra_pressure_levels: + path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/pressure_levels.zarr + dims: [time, x, y, pressure] + variables: + z: + pressure: + values: [1000,] + units: hPa + t: + pressure: + values: [800, ] + units: hPa + dim_mapping: + time: + method: rename + dim: time + state_feature: + method: stack_variables_by_var_name + dims: [pressure] + name_format: "{var_name}{pressure}m" + grid_index: + method: stack + dims: [x, y] + target_output_variable: state +""" -def modify_example_config_inputs_section(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 reading the config + Update provided config. Parameters ---------- - new_inputs_section: str - String with a new inputs section + config: str + String with config in yaml format + update: str + String with the update in yaml format Returns ------- config: Config - Modified config with the new inputs section replacing the old one - in the example 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 - - # Read the config - modified_config = mdp.Config.from_yaml(modified_yaml) + original_config = mdp.Config.from_yaml(config) + update = yaml.safe_load(update) + modified_config = original_config.to_dict() + modified_config.update(update) + modified_config = mdp.Config.from_dict(modified_config) return modified_config @pytest.mark.parametrize( - "new_inputs_section", + "base_config, new_inputs_section", [ - None, # Does not modify the example config - PRESSURE_LEVEL_TEST_SECTION, - HEIGHT_LEVEL_TEST_SECTION, - SINGLE_LEVEL_SELECTED_VARIABLES_TEST_SECTION, - SINGLE_LEVEL_DERIVED_VARIABLES_TEST_SECTION, + (BASE_CONFIG, "{}"), # Does not modify the example config + (BASE_CONFIG, PRESSURE_LEVEL_TEST_SECTION), + (BASE_CONFIG, HEIGHT_LEVEL_TEST_SECTION), + (BASE_CONFIG, SINGLE_LEVEL_SELECTED_VARIABLES_TEST_SECTION), + (BASE_CONFIG, SINGLE_LEVEL_DERIVED_VARIABLES_TEST_SECTION), ], ) -def test_selected_output_variables(new_inputs_section): +def test_selected_output_variables(base_config, new_inputs_section): """ Test that the variables specified in each input dataset are present in the output dataset. """ # Modify the example config - config = modify_example_config_inputs_section(new_inputs_section) + config = update_config(base_config, new_inputs_section) # Create the dataset ds = mdp.create_dataset(config=config) # Check that the output variables are the ones selected - for dataset_name, input_config in config.inputs.items(): + for _, input_config in config.inputs.items(): target_output_variable = input_config.target_output_variable # Get the expected selected variable names @@ -244,58 +255,26 @@ def test_selected_output_variables(new_inputs_section): pytest.fail(error_message) -INVALID_PRESSURE_LEVEL_TEST_SECTION = """\ -inputs: - danra_pressure_levels: - path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/pressure_levels.zarr - dims: [time, x, y, pressure] - variables: - z: - pressure: - values: [1000,] - units: hPa - t: - pressure: - values: [800, ] - units: hPa - dim_mapping: - time: - method: rename - dim: time - state_feature: - method: stack_variables_by_var_name - dims: [pressure] - name_format: "{var_name}{pressure}m" - grid_index: - method: stack - dims: [x, y] - target_output_variable: state -""" - - @pytest.mark.parametrize( - "new_inputs_section, expected_result", + "base_config, update, expected_result", [ ( - None, + BASE_CONFIG, + "{}", False, ), # Do not modify the example config - should return False since we're expecting no nans ( + BASE_CONFIG, INVALID_PRESSURE_LEVEL_TEST_SECTION, True, ), # Dataset with nans - should return True ], ) -def test_output_dataset_for_nans(new_inputs_section, expected_result): +def test_output_dataset_for_nans(base_config, update, expected_result): """ Test that the output dataset does not contain any nan values. """ - # Modify the example config - config = modify_example_config_inputs_section(new_inputs_section) - - # Create the dataset + config = update_config(base_config, update) ds = mdp.create_dataset(config=config) - - # Test that we have no nans - nan_in_ds = ds.isnull().any().compute().to_array().any().item() + nan_in_ds = any(ds.isnull().any().compute().to_array()) assert nan_in_ds == expected_result From 827bfb1a2f4b1e3a2cf9e3401ce6fffd2f5088f9 Mon Sep 17 00:00:00 2001 From: ealerskans Date: Fri, 14 Feb 2025 07:37:16 +0000 Subject: [PATCH 5/5] Update CHANGELOG --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2dcc56d..b582375 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixes +- fix bug which adds unwanted dimensions to the dataset [\#60](https://github.com/mllam/mllam-data-prep/pull/60), @ealerskans, @observingClouds - correct chunk size estimate [\#59](https://github.com/mllam/mllam-data-prep/pull/59), @ealerskans - fix bug arising when variables provided to derived functions are renamed [\#56](https://github.com/mllam/mllam-data-prep/pull/56), @leifdenby - ensure config fields defaulting to `None` are typed as `Optional` and fields defaulting to `{}` are given a default-factory so that serialization with default values works correctly [\#63](https://github.com/mllam/mllam-data-prep/pull/63), @leifdenby