diff --git a/docs/history.md b/docs/history.md index aa9dcd52..daabab95 100644 --- a/docs/history.md +++ b/docs/history.md @@ -1,14 +1,18 @@ # History +## Development Version + +* ADD: Added `apply_to_sweeps` function for applying custom operations to all sweeps in a `DataTree` radar volume Implemented by [@syedhamidali](https://github.com/syedhamidali), ({pull}`202`). + ## 0.6.5 (2024-09-20) -FIX: Azimuth dimension now labelled correctly for Halo Photonics data ({pull}`206`) by [@rcjackson](https://github.com/rcjackson). -FIX: Do not apply scale/offset in datamet reader, leave it to xarray instead ({pull}`209`) by [@kmuehlbauer](https://github.com/kmuehlbauer). +* FIX: Azimuth dimension now labelled correctly for Halo Photonics data ({pull}`206`) by [@rcjackson](https://github.com/rcjackson). +* FIX: do not apply scale/offset in datamet reader, leave it to xarray instead ({pull}`209`) by [@kmuehlbauer](https://github.com/kmuehlbauer). ## 0.6.4 (2024-08-30) -FIX: Notebooks are now conforming to ruff's style checks by [@rcjackson](https://github.com/rcjackson), ({pull}`199`) by [@rcjackson](https://github.com/rcjackson). -FIX: use dict.get() to retrieve attribute key and return "None" if not available, ({pull}`200`) by [@kmuehlbauer](https://github.com/kmuehlbauer) +* FIX: Notebooks are now conforming to ruff's style checks by [@rcjackson](https://github.com/rcjackson), ({pull}`199`) by [@rcjackson](https://github.com/rcjackson). +* FIX: use dict.get() to retrieve attribute key and return "None" if not available, ({pull}`200`) by [@kmuehlbauer](https://github.com/kmuehlbauer) ## 0.6.3 (2024-08-13) diff --git a/tests/test_util.py b/tests/test_util.py index 1552f5be..89e37cee 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -4,6 +4,7 @@ """Tests for `xradar` util package.""" +import datatree as dt import numpy as np import pytest import xarray as xr @@ -264,3 +265,115 @@ def test_get_sweep_keys(): dt["sneep_1"] = dt["sweep_1"] keys = util.get_sweep_keys(dt) assert keys == ["sweep_0", "sweep_1", "sweep_2", "sweep_3", "sweep_4", "sweep_5"] + + +def test_apply_to_sweeps(): + # Fetch the sample radar file + filename = DATASETS.fetch("sample_sgp_data.nc") + + # Open the radar file into a DataTree object + dtree = io.open_cfradial1_datatree(filename) + + # Define a simple function to test with apply_to_sweeps + def dummy_function(ds): + """A dummy function that adds a constant field to the dataset.""" + ds["dummy_field"] = ( + ds["reflectivity_horizontal"] * 0 + ) # Adding a field with all zeros + ds["dummy_field"].attrs = {"units": "dBZ", "long_name": "Dummy Field"} + return ds + + # Apply the dummy function to all sweeps using apply_to_sweeps + modified_dtree = util.apply_to_sweeps(dtree, dummy_function) + + # Verify that the dummy field has been added to each sweep + sweep_keys = util.get_sweep_keys(modified_dtree) + for key in sweep_keys: + assert ( + "dummy_field" in modified_dtree[key].data_vars + ), f"dummy_field not found in {key}" + assert modified_dtree[key]["dummy_field"].attrs["units"] == "dBZ" + assert modified_dtree[key]["dummy_field"].attrs["long_name"] == "Dummy Field" + + # Check that the original data has not been modified + assert ( + "dummy_field" not in dtree["/"].data_vars + ), "dummy_field should not be in the root node" + + # Test that an exception is raised when a function that causes an error is applied + with pytest.raises(ValueError, match="This is an intentional error"): + + def error_function(ds): + raise ValueError("This is an intentional error") + + util.apply_to_sweeps(dtree, error_function) + + +def test_apply_to_volume(): + # Fetch the sample radar file + filename = DATASETS.fetch("sample_sgp_data.nc") + + # Open the radar file into a DataTree object + dtree = io.open_cfradial1_datatree(filename) + + # Define a simple function to test with apply_to_volume + def dummy_function(ds): + """A dummy function that adds a constant field to the dataset.""" + ds["dummy_field"] = ( + ds["reflectivity_horizontal"] * 0 + ) # Adding a field with all zeros + ds["dummy_field"].attrs = {"units": "dBZ", "long_name": "Dummy Field"} + return ds + + # Apply the dummy function to all sweeps using apply_to_volume + modified_dtree = util.apply_to_volume(dtree, dummy_function) + + # Verify that the modified_dtree is an instance of DataTree + assert isinstance( + modified_dtree, dt.DataTree + ), "The result should be a DataTree instance." + + # Verify that the dummy field has been added to each sweep + sweep_keys = util.get_sweep_keys(modified_dtree) + for key in sweep_keys: + assert ( + "dummy_field" in modified_dtree[key].data_vars + ), f"dummy_field not found in {key}" + assert modified_dtree[key]["dummy_field"].attrs["units"] == "dBZ" + assert modified_dtree[key]["dummy_field"].attrs["long_name"] == "Dummy Field" + + # Check that the original DataTree (dtree) has not been modified + original_sweep_keys = util.get_sweep_keys(dtree) + for key in original_sweep_keys: + assert ( + "dummy_field" not in dtree[key].data_vars + ), f"dummy_field should not be in the original DataTree at {key}" + + # Test edge case: Apply a function that modifies only certain sweeps + def selective_function(ds): + """Only modifies sweeps with a specific condition.""" + if "reflectivity_horizontal" in ds: + ds["selective_field"] = ds["reflectivity_horizontal"] * 1 + return ds + + # Apply the selective function to all sweeps using apply_to_volume + selectively_modified_dtree = util.apply_to_volume(dtree, selective_function) + + # Verify that the selective field was added only where the condition was met + for key in sweep_keys: + if "reflectivity_horizontal" in modified_dtree[key].data_vars: + assert ( + "selective_field" in selectively_modified_dtree[key].data_vars + ), f"selective_field not found in {key} where it should have been added." + else: + assert ( + "selective_field" not in selectively_modified_dtree[key].data_vars + ), f"selective_field should not be present in {key}" + + # Test that an exception is raised when a function that causes an error is applied + with pytest.raises(ValueError, match="This is an intentional error"): + + def error_function(ds): + raise ValueError("This is an intentional error") + + util.apply_to_volume(dtree, error_function) diff --git a/xradar/io/export/cfradial1.py b/xradar/io/export/cfradial1.py index af93c660..1b5f7471 100644 --- a/xradar/io/export/cfradial1.py +++ b/xradar/io/export/cfradial1.py @@ -3,7 +3,6 @@ # Distributed under the MIT License. See LICENSE for more info. """ - CfRadial1 output ================ @@ -20,13 +19,14 @@ :toctree: generated/ {} - """ __all__ = [ "to_cfradial1", ] +__doc__ = __doc__.format("\n ".join(__all__)) + from importlib.metadata import version import numpy as np @@ -139,7 +139,6 @@ def _variable_mapper(dtree, dim0=None): combine_attrs="drop_conflicts", ) - # Check if specific variables exist before dropping them drop_variables = [ "sweep_fixed_angle", "sweep_number", @@ -277,7 +276,8 @@ def calculate_sweep_indices(dtree, dataset=None): def to_cfradial1(dtree=None, filename=None, calibs=True): """ Convert a radar datatree.DataTree to the CFRadial1 format - and save it to a file. + and save it to a file. Ensure that the resulting dataset + is well-formed and does not include specified extraneous variables. Parameters ---------- @@ -286,18 +286,19 @@ def to_cfradial1(dtree=None, filename=None, calibs=True): filename: str, optional The name of the output netCDF file. calibs: Bool, optional - calibration parameters + Whether to include calibration parameters. """ + # Generate the initial ds_cf using the existing mapping functions dataset = _variable_mapper(dtree) - # Check if radar_parameters, radar_calibration, and - # georeferencing_correction exist in dtree + # Handle calibration parameters if calibs: if "radar_calibration" in dtree: calib_params = dtree["radar_calibration"].to_dataset() calibs = _calib_mapper(calib_params) dataset.update(calibs) + # Add additional parameters if they exist in dtree if "radar_parameters" in dtree: radar_params = dtree["radar_parameters"].to_dataset() dataset.update(radar_params) @@ -306,8 +307,12 @@ def to_cfradial1(dtree=None, filename=None, calibs=True): radar_georef = dtree["georeferencing_correction"].to_dataset() dataset.update(radar_georef) - dataset.attrs = dtree.attrs + # Ensure that the data type of sweep_mode and similar variables matches + if "sweep_mode" in dataset.variables: + dataset["sweep_mode"] = dataset["sweep_mode"].astype("S") + # Update global attributes + dataset.attrs = dtree.attrs dataset.attrs["Conventions"] = "Cf/Radial" dataset.attrs["version"] = "1.2" xradar_version = version("xradar") diff --git a/xradar/util.py b/xradar/util.py index ea27430f..450531a3 100644 --- a/xradar/util.py +++ b/xradar/util.py @@ -23,6 +23,7 @@ "ipol_time", "rolling_dim", "get_sweep_keys", + "apply_to_sweeps", ] __doc__ = __doc__.format("\n ".join(__all__)) @@ -33,6 +34,7 @@ import io import warnings +import datatree as dt import numpy as np from scipy import interpolate @@ -522,3 +524,66 @@ def get_sweep_keys(dt): pass return sweep_group_keys + + +def apply_to_sweeps(dtree, func, *args, **kwargs): + """ + Applies a given function to all sweep nodes in the radar volume. + + Parameters + ---------- + dtree : DataTree + The DataTree object representing the radar volume. + func : function + The function to apply to each sweep. + *args : tuple + Additional positional arguments to pass to the function. + **kwargs : dict + Additional keyword arguments to pass to the function. + + Returns + ------- + DataTree + A new DataTree object with the function applied to all sweeps. + """ + # Create a new tree dictionary + tree = {"/": dtree.ds} # Start with the root Dataset + + # Add all nodes except the root + tree.update({node.path: node.ds for node in dtree.subtree if node.path != "/"}) + + # Apply the function to all sweep nodes and update the tree dictionary + tree.update( + { + node.path: func(dtree[node.path].to_dataset(), *args, **kwargs) + for node in dtree.match("sweep*").subtree + if node.path.startswith("/sweep") + } + ) + + # Return a new DataTree constructed from the modified tree dictionary + return dt.DataTree.from_dict(tree) + + +def apply_to_volume(dtree, func, *args, **kwargs): + """ + Alias for apply_to_sweeps. + Applies a given function to all sweep nodes in the radar volume. + + Parameters + ---------- + dtree : DataTree + The DataTree object representing the radar volume. + func : function + The function to apply to each sweep. + *args : tuple + Additional positional arguments to pass to the function. + **kwargs : dict + Additional keyword arguments to pass to the function. + + Returns + ------- + DataTree + A new DataTree object with the function applied to all sweeps. + """ + return apply_to_sweeps(dtree, func, *args, **kwargs)