Dtree constructor

docs/history.md

@@ -3,6 +3,7 @@
 ## 0.8.0 (2024-10-28)
 
 This is the first version which uses datatree directly from xarray. Thus, xarray is pinned to version >= 2024.10.0.
+* ENH: Refactoring all xradar backends to use `from_dict` datatree constructor. Test for `_get_required_root`, `_get_subgroup`, and `_get_radar_calibration` were also added ({pull}`221`) by [@aladinor](https://github.com/aladinor)
 * ENH: Added pytests to the missing functions in the `test_xradar` and `test_iris` in order to increase codecov in ({pull}`228`) by [@syedhamidali](https://github.com/syedhamidali).
 * ENH: Updated Readme ({pull}`226`) by [@syedhamidali](https://github.com/syedhamidali).
 * ADD: Added new module `transform` for transforming CF1 data to CF2 and vice versa ({pull}`224`) by [@syedhamidali](https://github.com/syedhamidali).

tests/io/test_io.py

@@ -109,7 +109,7 @@ def test_open_odim_datatree_sweep(odim_file, sweep):
         lswp = len([sweep])
     else:
         lswp = len(sweep)
-    assert len(dtree.groups[1:]) == lswp
+    assert len(dtree.match("sweep_*")) == lswp
 
 
 def test_open_odim_datatree(odim_file):
@@ -164,7 +164,7 @@ def test_open_odim_datatree(odim_file):
         200,
         200,
     ]
-    for i, grp in enumerate(dtree.groups[1:]):
+    for i, grp in enumerate(dtree.match("sweep_*")):
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
@@ -183,7 +183,7 @@ def test_open_odim_datatree(odim_file):
             "range",
         }
         assert np.round(ds.elevation.mean().values.item(), 1) == elevations[i]
-        assert ds.sweep_number.values == int(grp[7:])
+        assert ds.sweep_number.values == int(grp[6:])
 
 
 @pytest.mark.parametrize("first_dim", ["auto", "time"])
@@ -258,7 +258,7 @@ def test_open_gamic_datatree_sweep(gamic_file, sweep):
         lswp = len([sweep])
     else:
         lswp = len(sweep)
-    assert len(dtree.groups[1:]) == lswp
+    assert len(dtree.match("sweep_*")) == lswp
 
 
 def test_open_gamic_datatree(gamic_file):
@@ -319,7 +319,7 @@ def test_open_gamic_datatree(gamic_file):
         1000,
         1000,
     ]
-    for i, grp in enumerate(dtree.groups[1:]):
+    for i, grp in enumerate(dtree.match("sweep_*")):
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
@@ -545,7 +545,7 @@ def test_open_rainbow_datatree(rainbow_file):
     ]
     azimuths = [361] * 14
     ranges = [400] * 14
-    for i, grp in enumerate(dtree.groups[1:]):
+    for i, grp in enumerate(dtree.match("sweep_*")):
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
         assert set(ds.data_vars) & (
@@ -641,28 +641,27 @@ def test_open_iris_datatree(iris0_file):
     azimuths = [360] * 10
     ranges = [664] * 10
     i = 0
-    for grp in dtree.groups:
-        if grp.startswith("/sweep_"):
-            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
-            ) == set(moments)
-            assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
-                required_sweep_metadata_vars ^ {"azimuth", "elevation"}
-            )
-            assert set(ds.coords) == {
-                "azimuth",
-                "elevation",
-                "time",
-                "latitude",
-                "longitude",
-                "altitude",
-                "range",
-            }
-            assert np.round(ds.elevation.mean().values.item(), 1) == elevations[i]
-            assert ds.sweep_number == i
-            i += 1
+    for grp in dtree.match("sweep_*"):
+        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
+        ) == set(moments)
+        assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
+            required_sweep_metadata_vars ^ {"azimuth", "elevation"}
+        )
+        assert set(ds.coords) == {
+            "azimuth",
+            "elevation",
+            "time",
+            "latitude",
+            "longitude",
+            "altitude",
+            "range",
+        }
+        assert np.round(ds.elevation.mean().values.item(), 1) == elevations[i]
+        assert ds.sweep_number == i
+        i += 1
 
 
 def test_open_iris0_dataset(iris0_file):
@@ -879,36 +878,35 @@ def test_open_datamet_datatree(datamet_file):
     azimuths = [360] * 11
     ranges = [493, 493, 493, 664, 832, 832, 1000, 1000, 1332, 1332, 1332]
     i = 0
-    for grp in dtree.groups:
-        if grp.startswith("/sweep_"):
-            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
-            ) == set(moments)
-            assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
-                required_sweep_metadata_vars ^ {"azimuth", "elevation"}
-            )
-            assert set(ds.coords) == {
-                "azimuth",
-                "elevation",
-                "time",
-                "latitude",
-                "longitude",
-                "altitude",
-                "range",
-            }
-            assert np.round(ds.elevation.mean().values.item(), 1) == elevations[i]
-            assert ds.sweep_number == i
-            i += 1
+    for grp in dtree.match("sweep_*"):
+        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
+        ) == set(moments)
+        assert set(ds.data_vars) & (required_sweep_metadata_vars) == set(
+            required_sweep_metadata_vars ^ {"azimuth", "elevation"}
+        )
+        assert set(ds.coords) == {
+            "azimuth",
+            "elevation",
+            "time",
+            "latitude",
+            "longitude",
+            "altitude",
+            "range",
+        }
+        assert np.round(ds.elevation.mean().values.item(), 1) == elevations[i]
+        assert ds.sweep_number == i
+        i += 1
 
     # Try to reed single sweep
     dtree = open_datamet_datatree(datamet_file, sweep=1)
-    assert len(dtree.groups) == 2
+    assert len(dtree.groups) == 5
 
     # Try to read list of sweeps
     dtree = open_datamet_datatree(datamet_file, sweep=[1, 2])
-    assert len(dtree.groups) == 3
+    assert len(dtree.groups) == 6
 
 
 @pytest.mark.parametrize("first_dim", ["time", "auto"])
@@ -993,6 +991,7 @@ def test_cfradial_n_points_file(cfradial1n_file):
         assert ds.sweep_mode == "azimuth_surveillance"
 
 
+@pytest.mark.run(order=1)
 @pytest.mark.parametrize("sweep", ["sweep_0", 0, [0, 1], ["sweep_0", "sweep_1"]])
 @pytest.mark.parametrize(
     "nexradlevel2_files", ["nexradlevel2_gzfile", "nexradlevel2_bzfile"], indirect=True
@@ -1003,7 +1002,7 @@ def test_open_nexradlevel2_datatree_sweep(nexradlevel2_files, sweep):
         lswp = len([sweep])
     else:
         lswp = len(sweep)
-    assert len(dtree.groups[1:]) == lswp
+    assert len(dtree.match("sweep*")) == lswp
 
 
 @pytest.mark.parametrize(
@@ -1080,8 +1079,8 @@ def test_open_nexradlevel2_datatree(nexradlevel2_files):
         308,
         232,
     ]
-    assert len(dtree.groups[1:]) == 11
-    for i, grp in enumerate(dtree.groups[1:]):
+    assert len(dtree.groups[1:]) == 14
+    for i, grp in enumerate(dtree.match("sweep_*")):
         print(i)
         ds = dtree[grp].ds
         assert dict(ds.sizes) == {"azimuth": azimuths[i], "range": ranges[i]}
@@ -1101,4 +1100,4 @@ def test_open_nexradlevel2_datatree(nexradlevel2_files):
             "range",
         }
         assert np.round(ds.elevation.mean().values.item(), 1) == elevations[i]
-        assert ds.sweep_number.values == int(grp[7:])
+        assert ds.sweep_number.values == int(grp[6:])

tests/test_util.py

@@ -11,6 +11,11 @@
 
 import xradar as xd
 from xradar import io, model, util
+from xradar.io.backends.common import (
+    _get_radar_calibration,
+    _get_required_root_dataset,
+    _get_subgroup,
+)
 
 
 @pytest.fixture(
@@ -424,3 +429,57 @@ def dummy_function(ds, refl="none"):
         sweep_0.dummy_field.values
     )  # Convert NaNs to zero for comparison
     assert np.all(np.isclose(non_nan_values, 0))
+
+
+def test_get_required_root_dataset():
+
+    filename = DATASETS.fetch("cor-main131125105503.RAW2049")
+    sweeps = [f"sweep_{i}" for i in range(10)]
+    ls_ds = [xr.open_dataset(filename, engine="iris", group=sweep) for sweep in sweeps]
+    root = _get_required_root_dataset(ls_ds, optional=True)
+    elevations = [
+        0.5,
+        1.0,
+        2.0,
+        3.0,
+        5.0,
+        7.0,
+        10.0,
+        15.0,
+        20.0,
+        30.0,
+    ]
+    assert len(root.variables) == 10
+    assert root.variables["time_coverage_start"] == "2013-11-25T10:55:04Z"
+    assert root.variables["time_coverage_end"] == "2013-11-25T10:59:24Z"
+    np.testing.assert_equal(
+        root.variables["sweep_fixed_angle"].values, np.array(elevations)
+    )
+    assert len(list(root.attrs.keys())) == 10
+    assert root.attrs["instrument_name"] == "Corozal, Radar"
+    assert root.attrs["scan_name"] == "SURV_HV_300 "
+    assert root.attrs["comment"] == "AEROCIVIL OPERATIONAL DUAL POLE SCAN"
+
+
+def test_get_radar_calibration():
+    filename = DATASETS.fetch("DWD-Vol-2_99999_20180601054047_00.h5")
+    sweeps = [f"sweep_{i}" for i in range(10)]
+    ls_ds = [xr.open_dataset(filename, engine="gamic", group=sweep) for sweep in sweeps]
+    subgroup = _get_radar_calibration(ls_ds, model.radar_calibration_subgroup)
+    assert len(subgroup.variables) == 6
+    assert subgroup["noise_power_h"] == "-3.8298"
+    assert subgroup["rx_loss_h"] == "3"
+    assert subgroup["ant_gain_v"] == "43"
+    assert subgroup["ant_gain_h"] == "43"
+
+
+def test_get_subgroup():
+    filename = DATASETS.fetch("71_20181220_060628.pvol.h5")
+    sweeps = [f"sweep_{i}" for i in range(10)]
+    ls_ds = [xr.open_dataset(filename, engine="odim", group=sweep) for sweep in sweeps]
+    subgroup = _get_subgroup(ls_ds, model.radar_parameters_subgroup)
+    assert len(subgroup.variables) == 3
+    assert list(subgroup.variables) == ["longitude", "latitude", "altitude"]
+    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)

xradar/io/backends/cfradial1.py

@@ -33,7 +33,7 @@
 __doc__ = __doc__.format("\n   ".join(__all__))
 
 import numpy as np
-from xarray import DataTree, merge, open_dataset
+from xarray import Dataset, DataTree, merge, open_dataset
 from xarray.backends import NetCDF4DataStore
 from xarray.backends.common import BackendEntrypoint
 from xarray.backends.store import StoreBackendEntrypoint
@@ -49,7 +49,7 @@
     required_global_attrs,
     required_root_vars,
 )
-from .common import _maybe_decode
+from .common import _attach_sweep_groups, _maybe_decode
 
 
 def _get_required_root_dataset(ds, optional=True):
@@ -89,7 +89,6 @@ def _get_required_root_dataset(ds, optional=True):
     root.sweep_group_name.encoding["dtype"] = root.sweep_group_name.dtype
     # remove cf standard name
     root.sweep_group_name.attrs = []
-
     return root
 
 
@@ -297,6 +296,8 @@ def _get_radar_calibration(ds):
         subgroup = subgroup.rename_vars(calib_vars)
         subgroup.attrs = {}
         return subgroup
+    else:
+        return Dataset()
 
 
 def open_cfradial1_datatree(filename_or_obj, **kwargs):
@@ -354,24 +355,22 @@ def open_cfradial1_datatree(filename_or_obj, **kwargs):
         "/georeferencing_correction": _get_subgroup(
             ds, georeferencing_correction_subgroup
         ),
+        "/radar_calibration": _get_radar_calibration(ds),
     }
 
     # radar_calibration (connected with calib-dimension)
-    calib = _get_radar_calibration(ds)
-    if calib:
-        dtree["/radar_calibration"] = calib
-
-    sweep_child = list(
-        _get_sweep_groups(
-            ds,
-            sweep=sweep,
-            first_dim=first_dim,
-            optional=optional,
-            site_coords=site_coords,
-        ).values()
+    dtree = _attach_sweep_groups(
+        dtree,
+        list(
+            _get_sweep_groups(
+                ds,
+                sweep=sweep,
+                first_dim=first_dim,
+                optional=optional,
+                site_coords=site_coords,
+            ).values()
+        ),
     )
-    for i, sw in enumerate(sweep_child):
-        dtree[f"sweep_{i}"] = sw
     return DataTree.from_dict(dtree)