ADD: Add gridding support from the xradar object

pyart/xradar/accessor.py

@@ -5,16 +5,16 @@
 
 
 import copy
-from collections.abc import Hashable, Mapping
-from typing import Any, overload
 
 import numpy as np
 import pandas as pd
 from datatree import DataTree, formatting, formatting_html
 from datatree.treenode import NodePath
-from xarray import DataArray, Dataset, concat
+from xarray import concat
 from xarray.core import utils
 
+from ..core.transforms import antenna_vectors_to_cartesian
+
 
 class Xradar:
     def __init__(self, xradar, default_sweep="sweep_0"):
@@ -36,8 +36,15 @@ def __init__(self, xradar, default_sweep="sweep_0"):
         self.elevation = dict(data=self.combined_sweeps.elevation.values)
         self.fixed_angle = dict(data=self.combined_sweeps.sweep_fixed_angle.values)
         self.antenna_transition = None
-        self.latitude = dict(data=self.xradar["latitude"].values)
-        self.longitude = dict(data=self.xradar["longitude"].values)
+        self.latitude = dict(
+            data=np.expand_dims(self.xradar["latitude"].values, axis=0)
+        )
+        self.longitude = dict(
+            data=np.expand_dims(self.xradar["longitude"].values, axis=0)
+        )
+        self.altitude = dict(
+            data=np.expand_dims(self.xradar["altitude"].values, axis=0)
+        )
         self.sweep_end_ray_index = dict(
             data=self.combined_sweeps.ngates.groupby("sweep_number").max().values
         )
@@ -49,26 +56,16 @@ def __init__(self, xradar, default_sweep="sweep_0"):
         self.nrays = len(self.azimuth["data"])
         self.nsweeps = len(self.xradar.sweep_group_name)
         self.instrument_parameters = dict(**self.xradar["radar_parameters"].attrs)
+        self.init_gate_x_y_z()
+        self.init_gate_alt()
 
     def __repr__(self):
         return formatting.datatree_repr(self.xradar)
 
     def _repr_html_(self):
         return formatting_html.datatree_repr(self.xradar)
 
-    @overload
-    def __getitem__(self, key: Mapping) -> Dataset:  # type: ignore[misc]
-        ...
-
-    @overload
-    def __getitem__(self, key: Hashable) -> DataArray:  # type: ignore[misc]
-        ...
-
-    @overload
-    def __getitem__(self, key: Any) -> Dataset:
-        ...
-
-    def __getitem__(self: DataTree, key: str) -> DataTree | DataArray:
+    def __getitem__(self: DataTree, key):
         """
         Access child nodes, variables, or coordinates stored anywhere in this tree.
 
@@ -264,6 +261,40 @@ def get_gate_x_y_z(self, sweep, edges=False, filter_transitions=False):
         data = self.xradar[f"sweep_{sweep}"].xradar.georeference()
         return data["x"].values, data["y"].values, data["z"].values
 
+    def init_gate_x_y_z(self):
+        """Initialize or reset the gate_{x, y, z} attributes."""
+
+        ranges = self.range["data"]
+        azimuths = self.azimuth["data"]
+        elevations = self.elevation["data"]
+        cartesian_coords = antenna_vectors_to_cartesian(
+            ranges, azimuths, elevations, edges=False
+        )
+
+        if not hasattr(self, "gate_x"):
+            self.gate_x = dict()
+
+        if not hasattr(self, "gate_y"):
+            self.gate_y = dict()
+
+        if not hasattr(self, "gate_z"):
+            self.gate_z = dict()
+
+        self.gate_x = dict(data=cartesian_coords[0])
+        self.gate_y = dict(data=cartesian_coords[1])
+        self.gate_z = dict(data=cartesian_coords[2])
+
+    def init_gate_alt(self):
+        if not hasattr(self, "gate_altitude"):
+            self.gate_altitude = dict()
+
+        try:
+            self.gate_altitude = dict(data=self.altitude["data"] + self.gate_z["data"])
+        except ValueError:
+            self.gate_altitude = dict(
+                data=np.mean(self.altitude["data"]) + self.gate_z["data"]
+            )
+
     def _combine_sweeps(self, radar):
         # Loop through and extract the different datasets
         ds_list = []

tests/xradar/test_accessor.py

@@ -1,4 +1,6 @@
+import numpy as np
 import xradar as xd
+from numpy.testing import assert_allclose
 from open_radar_data import DATASETS
 
 import pyart
@@ -38,3 +40,21 @@ def test_add_field(filename=filename):
     radar.add_field("reflectivity", new_field)
     assert "reflectivity" in radar.fields
     assert radar["sweep_0"]["reflectivity"].shape == radar["sweep_0"]["DBZ"].shape
+
+
+def test_grid(filename=filename):
+    dtree = xd.io.open_cfradial1_datatree(
+        filename,
+        optional=False,
+    )
+    radar = pyart.xradar.Xradar(dtree)
+    grid = pyart.map.grid_from_radars(
+        (radar,),
+        grid_shape=(1, 11, 11),
+        grid_limits=((2000, 2000), (-100_000.0, 100_000.0), (-100_000.0, 100_000.0)),
+        fields=["DBZ"],
+    )
+    assert_allclose(grid.x["data"], np.arange(-100_000, 120_000, 20_000))
+    assert_allclose(
+        grid.fields["DBZ"]["data"][0, -1, 0], np.array(0.4243435), rtol=1e-03
+    )