Repository: openradar/xradar

Issue:
Issue with cfradial1 writer: ValueError: cannot rename 'sweep_fixed_angle' because it is not a variable or dimension in this dataset
* xradar version: 0.6.4
* Python version: 3.12.5
* Operating System: MacOS Sonoma 14.6.1

```python
import xradar as xd
import matplotlib.pyplot as plt
from open_radar_data import DATASETS

# Fetch a sample radar file
file = DATASETS.fetch('KATX20130717_195021_V06')

# Open the radar file into a DataTree object
dtree = xd.io.open_nexradlevel2_datatree(file)

xd.io.to_cfradial1(dtree, "out.nc")
```
```python
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/var/folders/yq/7n9vj93d6bq4gq64hx_6qgym0000gp/T/ipykernel_81303/1809764418.py in ?()
----> 1 xd.io.to_cfradial1(dtree)

~/Downloads/Others/xradar/xradar/io/export/cfradial1.py in ?(dtree, filename, calibs)
    287         The name of the output netCDF file.
    288     calibs: Bool, optional
    289         calibration parameters
    290     """
--> 291     dataset = _variable_mapper(dtree)
    292 
    293     # Check if radar_parameters, radar_calibration, and
    294     # georeferencing_correction exist in dtree

~/Downloads/Others/xradar/xradar/io/export/cfradial1.py in ?(dtree, dim0)
    101         Dataset containing mapped radar variables.
    102     """
    103 
    104     sweep_info = _sweep_info_mapper(dtree)
--> 105     vol_info = _main_info_mapper(dtree).drop_vars("fixed_angle")
    106     sweep_datasets = []
    107     for grp in dtree.groups:
    108         if "sweep" in grp:

~/Downloads/Others/xradar/xradar/io/export/cfradial1.py in ?(dtree)
     78     """
     79     dataset = (
     80         dtree.root.to_dataset()
     81         .drop_vars("sweep_group_name", errors="ignore")
---> 82         .rename({"sweep_fixed_angle": "fixed_angle"})
     83     )
     84     return dataset

~/miniconda3/envs/xradar-dev/lib/python3.12/site-packages/xarray/core/dataset.py in ?(self, name_dict, **names)
   4358         Dataset.rename_vars
   4359         Dataset.rename_dims
   4360         DataArray.rename
   4361         """
-> 4362         return self._rename(name_dict=name_dict, **names)

~/miniconda3/envs/xradar-dev/lib/python3.12/site-packages/xarray/core/dataset.py in ?(self, name_dict, **names)
   4296         """
   4297         name_dict = either_dict_or_kwargs(name_dict, names, "rename")
   4298         for k in name_dict.keys():
   4299             if k not in self and k not in self.dims:
-> 4300                 raise ValueError(
   4301                     f"cannot rename {k!r} because it is not a "
   4302                     "variable or dimension in this dataset"
   4303                 )

ValueError: cannot rename 'sweep_fixed_angle' because it is not a variable or dimension in this dataset

```

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/openradar/xradar/pull/202?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) Report
Attention: Patch coverage is `10.00000%` with `9 lines` in your changes missing coverage. Please review.
> Project coverage is 90.06%. Comparing base [(`34a9e82`)](https://app.codecov.io/gh/openradar/xradar/commit/34a9e8234189f921b1586001c84bdaf8f80f4c48?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) to head [(`f3c28b0`)](https://app.codecov.io/gh/openradar/xradar/commit/f3c28b086db88c332a33cc3a872ef3b797f256d0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar).

| [Files with missing lines](https://app.codecov.io/gh/openradar/xradar/pull/202?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) | Patch % | Lines |
|---|---|---|
| [xradar/util.py](https://app.codecov.io/gh/openradar/xradar/pull/202?src=pr&el=tree&filepath=xradar%2Futil.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar#diff-eHJhZGFyL3V0aWwucHk=) | 10.00% | [9 Missing :warning: ](https://app.codecov.io/gh/openradar/xradar/pull/202?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #202      +/-   ##
==========================================
- Coverage   91.98%   90.06%   -1.92%     
==========================================
  Files          23       23              
  Lines        4541     4551      +10     
==========================================
- Hits         4177     4099      -78     
- Misses        364      452      +88     
```

| [Flag](https://app.codecov.io/gh/openradar/xradar/pull/202/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) | Coverage Δ | |
|---|---|---|
| [notebooktests](https://app.codecov.io/gh/openradar/xradar/pull/202/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) | `?` | |
| [unittests](https://app.codecov.io/gh/openradar/xradar/pull/202/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) | `90.06% <10.00%> (-0.18%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/openradar/xradar/pull/202?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar).

kmuehlbauer: @syedhamidali Great idea!

We already have implemented the georeferencing to work over a tree. We might think to generalize that approach using this Pull Request.

For the georeferencing we used accessors and the `.pipe` method under the hood.

From my perspective it would make sense to implement your approach as accessor, too.

The API would be like this:

```python
dtree = dtree.xradar.apply(function, *args, **kwargs)
```

Where `apply` would be your `apply_to_sweep` and `function` the function to be applied. If done the right way with the accessors the function can be applied to DataArrays, Datasets and DataTrees.

Does that make sense?
syedhamidali: @kmuehlbauer Thanks!

Yeah, I agree that `dtree = dtree.xradar.apply(function, *args, **kwargs)` would be a great enhancement. It would make the approach more versatile and consistent with the existing API. However, I’m not yet familiar with how these accessors are built, so I’d appreciate any guidance or examples on how to implement this. Or, if you have the time to implement it yourself, that would be great too!
kmuehlbauer: @syedhamidali I'll respond with some more details later this weekend. Thanks for being supportive of the overall idea!
syedhamidali: @kmuehlbauer, when you get a chance, could you take a quick look at PR #203? I’d really appreciate your feedback.

Patch:
diff --git a/docs/history.md b/docs/history.md
index aa9dcd5..daabab9 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/xradar/io/export/cfradial1.py b/xradar/io/export/cfradial1.py
index af93c66..1b5f747 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 @@ Example::
    :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 ea27430..450531a 100644
--- a/xradar/util.py
+++ b/xradar/util.py
@@ -23,6 +23,7 @@ __all__ = [
     "ipol_time",
     "rolling_dim",
     "get_sweep_keys",
+    "apply_to_sweeps",
 ]
 
 __doc__ = __doc__.format("\n   ".join(__all__))
@@ -33,6 +34,7 @@ import importlib.util
 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)


Test patch:
diff --git a/tests/test_util.py b/tests/test_util.py
index 1552f5b..89e37ce 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)