Skip to content

Commit

Permalink
ADD: map_over_sweeps function
Browse files Browse the repository at this point in the history
  • Loading branch information
@syedhamidali
syedhamidali committed Oct 17, 2024
1 parent fb35f3e commit 3d46dbe
Show file tree
Hide file tree
Showing 7 changed files with 251 additions and 190 deletions.
2 changes: 1 addition & 1 deletion docs/usage.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ datamodel
importers
exporters
notebooks/Accessors
notebooks/Apply_Accessor
notebooks/Mapping_Sweeps
```

```{toctree}
Expand Down
71 changes: 52 additions & 19 deletions examples/notebooks/Apply_Accessor.ipynb → examples/notebooks/Mapping_Sweeps.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,16 +15,16 @@
"source": [
"Have you ever wondered how to efficiently apply operations to a full volume of radar sweeps, rather than processing each sweep individually? \n",
"\n",
"Xradar has the solution: the `apply` accessor. In this notebook, we’ll explore how you can leverage Xradar’s powerful `apply` functionality to perform volume-level operations on radar data.\n",
"Xradar has the solution: the `map_over_sweeps` accessor. In this notebook, we’ll explore how you can leverage Xradar’s powerful `map_over_sweeps` functionality to perform volume-level operations on radar data.\n",
"\n",
"In radar data analysis, it's common to work with multiple sweeps in a radar volume. Xradar allows you to apply custom functions across the entire dataset with ease, making complex operations, such as filtering reflectivity or calculating rain rate, efficient and scalable.\n",
"\n",
"Here's what you'll learn in this notebook:\n",
"- How to load and inspect radar data using Xradar’s `DataTree`.\n",
"- How to apply functions to process all radar sweeps in one go.\n",
"- How to load and inspect radar data using Xradar.\n",
"- How to apply functions to process all radar sweeps in one go using both conventional and decorator-based methods.\n",
"- How to visualize radar variables like reflectivity and rain rate before and after processing.\n",
"\n",
"Let’s dive into radar volume processing and learn how to apply operations across multiple sweeps efficiently."
"Let’s get into it!"
]
},
{
Expand Down Expand Up @@ -122,7 +122,7 @@
"id": "11",
"metadata": {},
"source": [
"### Example #1"
"## Example #1"
]
},
{
Expand All @@ -134,11 +134,13 @@
"source": [
"# It is just a demonstration, you can ignore the logic\n",
"def filter_radar(ds):\n",
" ds[\"DBZH_Filtered\"] = ds.where(\n",
" (ds[\"corrected_reflectivity_horizontal\"] > 10)\n",
" & (ds[\"corrected_reflectivity_horizontal\"] < 70)\n",
" & (ds[\"copol_coeff\"] > 0.8)\n",
" )[\"corrected_reflectivity_horizontal\"]\n",
" ds = ds.assign(\n",
" DBZH_Filtered=ds.where(\n",
" (ds[\"corrected_reflectivity_horizontal\"] > 10)\n",
" & (ds[\"corrected_reflectivity_horizontal\"] < 70)\n",
" & (ds[\"copol_coeff\"] > 0.85)\n",
" )[\"corrected_reflectivity_horizontal\"]\n",
" )\n",
" return ds"
]
},
Expand All @@ -150,15 +152,15 @@
"outputs": [],
"source": [
"# Apply the function across all sweeps\n",
"dtree = dtree.xradar.apply(filter_radar)"
"dtree = dtree.xradar.map_over_sweeps(filter_radar)"
]
},
{
"cell_type": "markdown",
"id": "14",
"metadata": {},
"source": [
"### Comparison"
"## Comparison"
]
},
{
Expand Down Expand Up @@ -234,7 +236,7 @@
"id": "17",
"metadata": {},
"source": [
"### Example #2"
"## Example #2"
]
},
{
Expand All @@ -250,7 +252,7 @@
" Z = 10.0 ** (dbz / 10.0)\n",
" return (Z / a) ** (1.0 / b)\n",
"\n",
" ds[\"RAIN_RATE\"] = _rain_rate(ds[ref_field])\n",
" ds = ds.assign(RAIN_RATE2=_rain_rate(ds[ref_field]))\n",
" ds[\"RAIN_RATE\"].attrs = {\"units\": \"mm/h\", \"long_name\": \"Rain Rate\"}\n",
" return ds"
]
Expand All @@ -263,7 +265,7 @@
"outputs": [],
"source": [
"# Apply the function across all sweeps\n",
"dtree = dtree.xradar.apply(calculate_rain_rate, ref_field=\"DBZH_Filtered\")\n",
"dtree = dtree.xradar.map_over_sweeps(calculate_rain_rate, ref_field=\"DBZH_Filtered\")\n",
"dtree"
]
},
Expand Down Expand Up @@ -332,22 +334,53 @@
"cell_type": "markdown",
"id": "22",
"metadata": {},
"source": [
"## With decorator"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "23",
"metadata": {},
"outputs": [],
"source": [
"@xd.map_over_sweeps\n",
"def calculate_rain_rate2(ds, ref_field=\"DBZH\"):\n",
" def _rain_rate(dbz, a=200.0, b=1.6):\n",
" Z = 10.0 ** (dbz / 10.0)\n",
" return (Z / a) ** (1.0 / b)\n",
"\n",
" ds = ds.assign(RAIN_RATE2=_rain_rate(ds[ref_field]))\n",
" ds.RAIN_RATE.attrs = {\"units\": \"mm/h\", \"long_name\": \"Rain Rate\"}\n",
" return ds\n",
"\n",
"\n",
"# invocation via decorator + pipe\n",
"dtree3 = dtree.pipe(calculate_rain_rate2, ref_field=\"DBZH_Filtered\")\n",
"display(dtree3[\"sweep_0\"])"
]
},
{
"cell_type": "markdown",
"id": "24",
"metadata": {},
"source": [
"## Conclusion"
]
},
{
"cell_type": "markdown",
"id": "23",
"id": "25",
"metadata": {},
"source": [
"In this notebook, we demonstrated how to handle and process radar data efficiently using Xradar’s `apply` accessor. By applying operations across an entire volume of radar sweeps, you can streamline your workflow and avoid the need to manually process each sweep.\n",
"In this notebook, we demonstrated how to handle and process radar data efficiently using Xradar’s `map_over_sweeps` accessor. By applying operations across an entire volume of radar sweeps, you can streamline your workflow and avoid the need to manually process each sweep.\n",
"\n",
"We explored two key use cases:\n",
"- **Filtering Reflectivity**: We applied a custom filtering function across all sweeps in the radar dataset, allowing us to isolate meaningful reflectivity values based on specific criteria.\n",
"- **Calculating Rain Rate**: Using the radar reflectivity data, we calculated the rain rate for each sweep, demonstrating how to perform scientific computations across multiple sweeps with minimal effort.\n",
"\n",
"The `apply` functionality in Xradar opens the door to performing various radar data processing tasks efficiently. Whether it's filtering, calculating derived quantities like rain rate, or applying more complex algorithms, Xradar simplifies working with radar volumes, making it easier to scale your analysis."
"The `map_over_sweeps` functionality in Xradar opens the door to performing various radar data processing tasks efficiently. Whether it's filtering, calculating derived quantities like rain rate, or applying more complex algorithms, Xradar simplifies working with radar volumes, making it easier to scale your analysis."
]
}
],
Expand All @@ -362,7 +395,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.5"
"version": "3.12.7"
}
},
"nbformat": 4,
Expand Down
132 changes: 54 additions & 78 deletions tests/test_accessors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,9 @@
# Copyright (c) 2022-2023, openradar developers.
# Distributed under the MIT License. See LICENSE for more info.

import datatree as dt
import numpy as np
import pytest
from numpy.testing import assert_almost_equal
import xarray as xr
from numpy.testing import assert_almost_equal, assert_raises
from open_radar_data import DATASETS

import xradar as xd
Expand Down Expand Up @@ -39,7 +38,7 @@ def test_georeference_dataset():

def test_georeference_datatree():
radar = xd.model.create_sweep_dataset()
tree = dt.DataTree.from_dict({"sweep_0": radar})
tree = xr.DataTree.from_dict({"sweep_0": radar})
geo = tree.xradar.georeference()["sweep_0"]
assert_almost_equal(
geo["x"].values[:3, 0], np.array([0.436241, 1.3085901, 2.1805407])
Expand All @@ -52,103 +51,80 @@ def test_georeference_datatree():
)


def test_xradar_datatree_accessor_apply():
def test_map_over_sweeps_apply_dummy_function():
"""
Test applying a dummy function to all sweep nodes using map_over_sweeps.
"""
# Fetch the sample radar file
filename = DATASETS.fetch("sample_sgp_data.nc")

# Open the radar file into a DataTree object
dtree = xd.io.open_cfradial1_datatree(filename)

# Define a simple function to test with the apply method
# Define a simple dummy function that adds a constant field to the dataset
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 = {
"unit": "dBZ",
"long_name": "Dummy Field",
}
ds = ds.assign(
dummy_field=ds["reflectivity_horizontal"] * 0
) # Field with zeros
ds["dummy_field"].attrs = {"unit": "dBZ", "long_name": "Dummy Field"}
return ds

# Apply the dummy function using the xradar accessor
dtree = dtree.xradar.apply(dummy_function)

# Verify that the dummy field has been added to each sweep
for key in xd.util.get_sweep_keys(dtree):
assert "dummy_field" in dtree[key].data_vars, f"dummy_field not found in {key}"
assert dtree[key]["dummy_field"].attrs["unit"] == "dBZ"
assert dtree[key]["dummy_field"].attrs["long_name"] == "Dummy Field"
# Apply using map_over_sweeps accessor
dtree_modified = dtree.xradar.map_over_sweeps(dummy_function)

# Test that the original ValueError is raised when a function that causes an error is applied
with pytest.raises(ValueError, match="This is an intentional error"):
# Check that the new field exists in sweep_0 and has the correct attributes
sweep_0 = dtree_modified["sweep_0"]
assert "dummy_field" in sweep_0.data_vars
assert sweep_0.dummy_field.attrs["unit"] == "dBZ"
assert sweep_0.dummy_field.attrs["long_name"] == "Dummy Field"

def error_function(ds):
raise ValueError("This is an intentional error")
# Ensure all non-NaN values are 0 (accounting for -0.0 and NaN values)
non_nan_values = np.nan_to_num(
sweep_0.dummy_field.values
) # Convert NaNs to zero for comparison
assert np.all(np.isclose(non_nan_values, 0))

dtree.xradar.apply(error_function)


def test_xradar_dataset_accessor_apply():
def test_map_over_sweeps_non_sweep_nodes():
"""
Test that non-sweep nodes remain unchanged when using map_over_sweeps.
"""
# Fetch the sample radar file
filename = DATASETS.fetch("sample_sgp_data.nc")

# Open the radar file into a DataTree object and extract a Dataset
# Open the radar file into a DataTree object and add a non-sweep node
dtree = xd.io.open_cfradial1_datatree(filename)
ds = dtree["sweep_0"].to_dataset() # Extracting the Dataset from one sweep
non_sweep_data = xr.Dataset({"non_sweep_data": ("dim", np.arange(10))})
dtree["non_sweep_node"] = non_sweep_data

# Define a simple function to test with the apply method
# Define a simple function that only modifies sweep nodes
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 = {"unit": "dBZ", "long_name": "Dummy Field"}
if "range" in ds.dims:
ds = ds.assign(
dummy_field=ds["reflectivity_horizontal"] * 0
) # Field with zeros
return ds

# Apply the dummy function using the xradar accessor
ds = ds.xradar.apply(dummy_function)

# Verify that the dummy field has been added
assert "dummy_field" in ds.data_vars, "dummy_field not found in dataset"
assert ds["dummy_field"].attrs["unit"] == "dBZ"
assert ds["dummy_field"].attrs["long_name"] == "Dummy Field"

# Test that the original ValueError 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")

ds.xradar.apply(error_function)
# Apply using map_over_sweeps
dtree_modified = dtree.xradar.map_over_sweeps(dummy_function)

# Check that non-sweep nodes remain unchanged
assert "non_sweep_data" in dtree_modified["non_sweep_node"].data_vars
assert "dummy_field" not in dtree_modified["non_sweep_node"].data_vars
assert "dummy_field" in dtree_modified["sweep_0"].data_vars

def test_xradar_dataarray_accessor_apply():
# Fetch the sample radar file
filename = DATASETS.fetch("sample_sgp_data.nc")

# Open the radar file into a DataTree object, extract a Dataset, and then a DataArray
dtree = xd.io.open_cfradial1_datatree(filename)
ds = dtree["sweep_0"].to_dataset() # Extracting the Dataset from one sweep
da = ds["reflectivity_horizontal"] # Extracting a DataArray from the Dataset

# Define a simple function to test with the apply method
def dummy_function(da):
"""A dummy function that adds a constant value to the data."""
return da + 10 # Add 10 to every element in the DataArray

# Apply the dummy function using the xradar accessor
da_modified = da.xradar.apply(dummy_function)

# Verify that the data was modified correctly using numpy's allclose
assert np.allclose(
da_modified.values, ds["reflectivity_horizontal"].values + 10, equal_nan=True
), "DataArray values not correctly modified"

# Test that the original ValueError is raised when a function that causes an error is applied
with pytest.raises(ValueError, match="This is an intentional error"):
def test_map_over_sweeps_invalid_input():
"""
Test map_over_sweeps with invalid input to ensure appropriate error handling.
"""
radar = xd.model.create_sweep_dataset()
tree = xr.DataTree.from_dict({"sweep_0": radar})

def error_function(da):
raise ValueError("This is an intentional error")
# Define a function that raises an error for invalid input
def invalid_rain_rate(ds, ref_field="INVALID_FIELD"):
return ds[ref_field] # This will raise a KeyError if the field is not present

da.xradar.apply(error_function)
# Expect a KeyError when applying the function with an invalid field
with assert_raises(KeyError):
tree.xradar.map_over_sweeps(invalid_rain_rate, ref_field="INVALID_FIELD")
Loading

0 comments on commit 3d46dbe

Please sign in to comment.