diff --git a/pyart/xradar/accessor.py b/pyart/xradar/accessor.py index 00f4584830..2f236b799f 100644 --- a/pyart/xradar/accessor.py +++ b/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,6 +56,8 @@ 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) @@ -56,19 +65,7 @@ def __repr__(self): 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 = [] diff --git a/tests/xradar/test_accessor.py b/tests/xradar/test_accessor.py index 282af3cd54..5cac06f0ab 100644 --- a/tests/xradar/test_accessor.py +++ b/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 + )