ADD: function to select dataset variables in sweep

ADD: function to get dataset variables in sweep FIX: typo in accessors module: Dataarray -> Dataset
openradar · Jan 28, 2025 · 17bdded · 17bdded
1 parent 830d86b
commit 17bdded
Show file tree

Hide file tree

Showing 6 changed files with 119 additions and 36 deletions.
diff --git a/docs/datamodel.md b/docs/datamodel.md
@@ -62,7 +62,7 @@ Internal Representation: {py:class}`xarray:xarray.DataArray`
 
 - DBZH, radar_equivalent_reflectivity_factor_h
 - DBZV, radar_equivalent_reflectivity_factor_v
-- and many more, see {class}`xradar.model.sweep_dataset_vars`
+- and many more, see {class}`xradar.model.sweep_observed_vars`
 
 Internal Representation: {py:class}`xarray:xarray.DataArray`
 

diff --git a/tests/io/test_io.py b/tests/io/test_io.py
@@ -24,9 +24,9 @@
     open_rainbow_datatree,
 )
 from xradar.model import (
-    non_standard_sweep_dataset_vars,
+    non_standard_sweep_observed_vars,
     required_sweep_metadata_vars,
-    sweep_dataset_vars,
+    sweep_observed_vars,
 )
 
 
@@ -69,7 +69,7 @@ def test_open_cfradial1_datatree(cfradial1_file):
         assert i == int(grp[7:])
         assert dict(ds.sizes) == {"time": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == set(moments)
         assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
             required_sweep_metadata_vars ^ {"azimuth", "elevation"}
@@ -89,15 +89,15 @@ def test_open_cfradial1_dataset(cfradial1_file):
     with xr.open_dataset(cfradial1_file, group="sweep_0", engine="cfradial1") as ds:
         assert list(ds.dims) == ["azimuth", "range"]
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {"DBZ", "VR"}
         assert ds.sweep_number == 0
 
     # open last sweep group
     with xr.open_dataset(cfradial1_file, group="sweep_8", engine="cfradial1") as ds:
         assert list(ds.dims) == ["azimuth", "range"]
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {"DBZ", "VR"}
         assert ds.sweep_number == 8
 
@@ -168,7 +168,7 @@ def test_open_odim_datatree(odim_file):
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == set(moments)
         assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
             required_sweep_metadata_vars ^ {"azimuth", "elevation"}
@@ -200,7 +200,7 @@ def test_open_odim_dataset(odim_file, first_dim, fix_second_angle):
         dim0 = "time" if first_dim == "time" else "azimuth"
         assert dict(ds.sizes) == {dim0: 360, "range": 1200}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {"WRADH", "VRADH", "PHIDP", "DBZH", "RHOHV", "KDP", "TH", "ZDR"}
         assert ds.sweep_number == 0
 
@@ -214,7 +214,7 @@ def test_open_odim_dataset(odim_file, first_dim, fix_second_angle):
     ) as ds:
         assert dict(ds.sizes) == {dim0: 360, "range": 280}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {"VRADH", "KDP", "WRADH", "TH", "RHOHV", "PHIDP", "ZDR", "DBZH"}
         assert ds.sweep_number == 11
 
@@ -323,7 +323,7 @@ def test_open_gamic_datatree(gamic_file):
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == set(moments)
         assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
             required_sweep_metadata_vars ^ {"azimuth", "elevation"}
@@ -355,7 +355,7 @@ def test_open_gamic_dataset(gamic_file, first_dim, fix_second_angle):
         dim0 = "time" if first_dim == "time" else "azimuth"
         assert dict(ds.sizes) == {dim0: 361, "range": 360}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {
             "WRADH",
             "WRADV",
@@ -382,7 +382,7 @@ def test_open_gamic_dataset(gamic_file, first_dim, fix_second_angle):
     ) as ds:
         assert dict(ds.sizes) == {dim0: 360, "range": 1000}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {
             "WRADH",
             "WRADV",
@@ -437,7 +437,7 @@ def test_open_furuno_scn_dataset(furuno_scn_file):
     with xr.open_dataset(furuno_scn_file, first_dim="time", engine="furuno") as ds:
         assert dict(ds.sizes) == {"time": 1376, "range": 602}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {"KDP", "VRADH", "ZDR", "DBZH", "WRADH", "RHOHV", "PHIDP"}
         for key, value in ds.data_vars.items():
             if key in [
@@ -473,7 +473,7 @@ def test_open_furuno_scnx_dataset(furuno_scnx_file):
     with xr.open_dataset(furuno_scnx_file, first_dim="time", engine="furuno") as ds:
         assert dict(ds.sizes) == {"time": 722, "range": 936}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {"KDP", "VRADH", "ZDR", "DBZH", "WRADH", "RHOHV", "PHIDP"}
 
         for key, value in ds.data_vars.items():
@@ -549,7 +549,7 @@ def test_open_rainbow_datatree(rainbow_file):
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == set(moments)
         assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
             required_sweep_metadata_vars ^ {"azimuth", "elevation"}
@@ -572,7 +572,7 @@ def test_open_rainbow_dataset(rainbow_file):
     with xr.open_dataset(rainbow_file, group="sweep_0", engine="rainbow") as ds:
         assert dict(ds.sizes) == {"azimuth": 361, "range": 400}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {
             "DBZH",
         }
@@ -582,7 +582,7 @@ def test_open_rainbow_dataset(rainbow_file):
     with xr.open_dataset(rainbow_file, group="sweep_13", engine="rainbow") as ds:
         assert dict(ds.sizes) == {"azimuth": 361, "range": 400}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {
             "DBZH",
         }
@@ -645,7 +645,7 @@ def test_open_iris_datatree(iris0_file):
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == set(moments)
         assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
             required_sweep_metadata_vars ^ {"azimuth", "elevation"}
@@ -669,7 +669,7 @@ def test_open_iris0_dataset(iris0_file):
     with xr.open_dataset(iris0_file, group="sweep_0", engine="iris") as ds:
         assert dict(ds.sizes) == {"azimuth": 360, "range": 664}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {
             "DBZH",
             "VRADH",
@@ -684,7 +684,7 @@ def test_open_iris0_dataset(iris0_file):
     with xr.open_dataset(iris0_file, group="sweep_9", engine="iris") as ds:
         assert dict(ds.sizes) == {"azimuth": 360, "range": 664}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {
             "DBZH",
             "VRADH",
@@ -711,7 +711,7 @@ def test_open_iris1_dataset(iris1_file):
     with xr.open_dataset(iris1_file, group="sweep_0", engine="iris") as ds:
         assert dict(ds.sizes) == {"azimuth": 359, "range": 833}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {
             "DBZH",
             "KDP",
@@ -824,7 +824,7 @@ def test_open_datamet_dataset(datamet_file):
     ) as ds:
         assert dict(ds.sizes) == {"azimuth": 360, "range": 493}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {"DBTH", "DBZH", "KDP", "PHIDP", "RHOHV", "VRADH", "WRADH", "ZDR"}
         assert ds.sweep_number == 0
 
@@ -836,7 +836,7 @@ def test_open_datamet_dataset(datamet_file):
     ) as ds:
         assert dict(ds.sizes) == {"azimuth": 360, "range": 1332}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == {"DBTH", "DBZH", "KDP", "PHIDP", "RHOHV", "VRADH", "WRADH", "ZDR"}
         assert ds.sweep_number == 10
 
@@ -882,7 +882,7 @@ def test_open_datamet_datatree(datamet_file):
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == set(moments)
         assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
             required_sweep_metadata_vars ^ {"azimuth", "elevation"}
@@ -975,7 +975,7 @@ def test_cfradial_n_points_file(cfradial1n_file):
             "frequency": 1,
         }
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == set(moments)
         assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
             required_sweep_metadata_vars ^ {"azimuth", "elevation"}
@@ -1085,7 +1085,7 @@ def test_open_nexradlevel2_datatree(nexradlevel2_files):
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
-            sweep_dataset_vars | non_standard_sweep_dataset_vars
+            sweep_observed_vars | non_standard_sweep_observed_vars
         ) == set(moments[i])
         assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
             required_sweep_metadata_vars ^ {"azimuth", "elevation"}

diff --git a/tests/test_util.py b/tests/test_util.py
@@ -483,3 +483,30 @@ def test_get_subgroup():
     np.testing.assert_almost_equal(subgroup.longitude.values.item(), 151.20899963378906)
     np.testing.assert_almost_equal(subgroup.latitude.values.item(), -33.700801849365234)
     assert isinstance(subgroup.altitude.values.item(), float)
+
+
+def test_get_observed_variables():
+
+    filename = DATASETS.fetch("cor-main131125105503.RAW2049")
+    sweep = xr.open_dataset(filename, group="sweep_7", engine="iris")
+    dataset_vars = util.get_dataset_variables(sweep)
+
+    assert dataset_vars == ["DBZH", "VRADH", "ZDR", "KDP", "PHIDP", "RHOHV"]
+
+
+def test_select_observed_variables():
+
+    filename = DATASETS.fetch("DWD-Vol-2_99999_20180601054047_00.h5")
+    sweep = xr.open_dataset(filename, group="sweep_7", engine="gamic")
+    sweep = sweep.assign(quality1=sweep["DBZH"] < 60)
+    sweep["DBZH"].attrs["ancillary_variables"] = ["quality1"]
+    sweep = sweep.assign(quality2=sweep["RHOHV"] < 0.5)
+    sweep["RHOHV"].attrs["ancillary_variables"] = ["quality2"]
+
+    select_vars = ["DBZH", "PHIDP"]
+    sweep = util.select_dataset_variables(sweep, variables=select_vars)
+    metadata = list(model.required_sweep_metadata_vars)
+    metadata.remove("elevation")
+    metadata.remove("azimuth")
+
+    assert set(sweep.data_vars) == set(select_vars + ["quality1"] + metadata)
diff --git a/xradar/accessors.py b/xradar/accessors.py
@@ -120,7 +120,7 @@ def get_crs(self):
 
 @xr.register_dataset_accessor("xradar")
 class XradarDataSetAccessor(XradarAccessor):
-    """Adds a number of xradar specific methods to xarray.DataArray objects."""
+    """Adds a number of xradar specific methods to xarray.Dataset objects."""
 
     def georeference(
         self, earth_radius=None, effective_radius_fraction=None

diff --git a/xradar/model.py b/xradar/model.py
@@ -50,8 +50,8 @@
     "sweep_coordinate_vars",
     "required_sweep_metadata_vars",
     "optional_sweep_metadata_vars",
-    "sweep_dataset_vars",
-    "non_standard_sweep_dataset_vars",
+    "sweep_observed_vars",
+    "non_standard_sweep_observed_vars",
     "determine_cfradial2_sweep_variables",
     "conform_cfradial2_sweep_group",
 ]
@@ -172,8 +172,8 @@
     "n_samples",
 }
 
-#: sweep dataset variable names
-sweep_dataset_vars = {
+#: sweep observed variable names
+sweep_observed_vars = {
     "DBZH",
     "DBZV",
     "ZH",
@@ -213,8 +213,8 @@
     "REC",
 }
 
-#: non-standard sweep dataset variable names
-non_standard_sweep_dataset_vars = {
+#: non-standard sweep observed variable names
+non_standard_sweep_observed_vars = {
     "DBZ",
     "VEL",
     "VR",
@@ -373,7 +373,7 @@
 #: required moment attributes
 moment_attrs = {"standard_name", "long_name", "units"}
 
-# todo: align this with sweep_dataset_vars
+# todo: align this with sweep_observed_vars
 #: CfRadial 2.1 / FM301 / ODIM_H5 mapping
 sweep_vars_mapping = {
     "DBZH": {