ADD: New function - Apply to sweeps

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)
 

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)

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")

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)