adding radar parameters to xradar iris datatree

AUTHORS.md

@@ -10,3 +10,4 @@
 
 * Edouard Goudenhoofdt <[email protected]>
 * Hamid Ali Syed <[email protected]>
+* Alfonso Ladino <[email protected]>

docs/history.md

@@ -2,6 +2,8 @@
 
 ## Development Version (unreleased)
 
+* ENH: Adding global variables and attributes to iris datatree ({pull}`166`) by [@aladinor](https://github.com/aladinor)
+
 ## 0.5.0 (2024-03-28)
 
 * MNT: Update GitHub actions, address DeprecationWarnings ({pull}`153`) by [@kmuehlbauer](https://github.com/kmuehlbauer).

examples/notebooks/multiple-sweeps-into-volume-scan.ipynb

@@ -144,9 +144,10 @@
     "    sweeps = list(i.children.keys())\n",
     "    print(f\"task sweeps: {sweeps}\")\n",
     "    for j in sweeps:\n",
-    "        print(\n",
-    "            f\"{j}: {i[j].sweep_fixed_angle.values: .1f} [deg], {i[j].range.values[-1] / 1e3:.1f} [km]\"\n",
-    "        )\n",
+    "        if j.startswith(\"sweep\"):\n",
+    "            print(\n",
+    "                f\"{j}: {i[j].sweep_fixed_angle.values: .1f} [deg], {i[j].range.values[-1] / 1e3:.1f} [km]\"\n",
+    "            )\n",
     "    print(\"----------------------------------------------------------------\")"
    ]
   },
@@ -437,9 +438,22 @@
    "source": [
     "display(vcps_back)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "96b5c030-3087-402b-b33b-6daf6d6e6214",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
@@ -450,7 +464,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16"
+   "version": "3.11.6"
   }
  },
  "nbformat": 4,

tests/io/test_io.py

@@ -616,26 +616,29 @@ def test_open_iris_datatree(iris0_file):
     ]
     azimuths = [360] * 10
     ranges = [664] * 10
-    for i, grp in enumerate(dtree.groups[1:]):
-        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 = 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
 
 
 def test_open_iris0_dataset(iris0_file):

xradar/io/backends/iris.py

@@ -60,6 +60,11 @@
     get_longitude_attrs,
     get_range_attrs,
     moment_attrs,
+    optional_root_attrs,
+    optional_root_vars,
+    radar_parameters_subgroup,
+    required_global_attrs,
+    required_root_vars,
     sweep_vars_mapping,
 )
 from .common import _assign_root, _attach_sweep_groups
@@ -981,7 +986,6 @@
     ]
 )
 
-
 # status_antenna_info Structure
 # 4.3.40, page 51
 STATUS_ANTENNA_INFO = OrderedDict(
@@ -1697,7 +1701,6 @@
     ]
 )
 
-
 # some length's of data structures
 LEN_STRUCTURE_HEADER = struct.calcsize(_get_fmt_string(STRUCTURE_HEADER))
 LEN_PRODUCT_HDR = struct.calcsize(_get_fmt_string(PRODUCT_HDR))
@@ -2466,7 +2469,6 @@
     ]
 )
 
-
 PRODUCT_DATA_TYPE_CODES = OrderedDict(
     [
         (0, {"name": "NULL", "struct": SPARE_PSI_STRUCT}),
@@ -2529,7 +2531,6 @@
     ]
 )
 
-
 RECORD_BYTES = 6144
 
 
@@ -3951,10 +3952,75 @@
     sid, opener = _check_iris_file(filename)
     with opener(filename, loaddata=False) as ds:
         # make this CfRadial2 conform
-        keys = [f"sweep_{i-1}" for i in list(ds.data.keys())]
+        keys = [f"sweep_{i - 1}" for i in list(ds.data.keys())]
     return keys
 
 
+def _get_required_root_dataset(ls_ds, optional=True):
+    """Extract Root Dataset."""
+    # keep only defined mandatory and defined optional variables per default
+    data_var = set(
+        [x for xs in [sweep.variables.keys() for sweep in ls_ds] for x in xs]
+    )
+    remove_root = set(data_var) ^ set(required_root_vars)
+    if optional:
+        remove_root ^= set(optional_root_vars)
+    remove_root ^= {
+        "fixed_angle",
+        "sweep_group_name",
+        "sweep_fixed_angle",
+    }
+    remove_root &= data_var
+    root = [sweep.drop_vars(remove_root) for sweep in ls_ds]
+    root_vars = set(
+        [x for xs in [sweep.variables.keys() for sweep in root] for x in xs]
+    )
+    # rename variables
+    # todo: find a more easy method not iterating over all variables
+    for k in root_vars:
+        rename = optional_root_vars.get(k, None)
+        if rename:
+            root = [sweep.rename_vars({k: rename}) for sweep in root]
+
+    root_vars = set(
+        [x for xs in [sweep.variables.keys() for sweep in root] for x in xs]
+    )
+    ds_vars = [sweep[root_vars] for sweep in ls_ds]
+
+    root = xr.concat(ds_vars, dim="sweep").reset_coords()
+    # keep only defined mandatory and defined optional attributes per default
+    attrs = root.attrs.keys()
+    remove_attrs = set(attrs) ^ set(required_global_attrs)
+    if optional:
+        remove_attrs ^= set(optional_root_attrs)
+    for k in remove_attrs:
+        root.attrs.pop(k, None)
+    # creating a copy of the dataset list for using the _assing_root function.
+    # and get the variabes/attributes for the root dataset
+    ls = ls_ds.copy()
+    ls.insert(0, xr.Dataset())
+    dtree = DataTree(data=_assign_root(ls), name="root")
+    root = root.assign(dtree.variables)
+    root.attrs = dtree.attrs
+    return root
+
+
+def _get_subgroup(ls_ds: list[xr.Dataset], subdict):
+    """Get iris-sigmet root metadata group.
+    Variables are fetched from the provided Dataset according to the subdict dictionary.
+    """
+    meta_vars = subdict
+    data_vars = set([x for xs in [ds.variables.keys() for ds in ls_ds] for x in xs])
+    extract_vars = set(data_vars) & set(meta_vars)
+    subgroup = xr.concat([ds[extract_vars] for ds in ls_ds], "sweep")
+    for k in subgroup.data_vars:
+        rename = meta_vars[k]
+        if rename:
+            subgroup = subgroup.rename_vars({k: rename})
+    subgroup.attrs = {}
+    return subgroup
+
+
 class IrisBackendEntrypoint(BackendEntrypoint):
     """Xarray BackendEntrypoint for IRIS/Sigmet data."""
 
@@ -4092,14 +4158,17 @@
     else:
         sweeps = _get_iris_group_names(filename_or_obj)
 
-    ds = [
+    ls_ds: list[xr.Dataset] = [
         xr.open_dataset(filename_or_obj, group=swp, engine="iris", **kwargs)
         for swp in sweeps
     ]
-
-    ds.insert(0, xr.Dataset())
-
+    # get the datatree root
+    root = _get_required_root_dataset(ls_ds)
     # create datatree root node with required data
-    dtree = DataTree(data=_assign_root(ds), name="root")
-    # return datatree with attached sweep child nodes
-    return _attach_sweep_groups(dtree, ds[1:])
+    dtree = DataTree(data=root, name="root")
+    # get radar_parameters group
+    subgroup = _get_subgroup(ls_ds, radar_parameters_subgroup)
+    # attach radar_parameter group
+    DataTree(subgroup, name="radar_parameters", parent=dtree)
+    # return Datatree attaching the sweep child nodes
+    return _attach_sweep_groups(dtree, ls_ds)