diff --git a/docs/history.md b/docs/history.md index a6e17616..abdc1312 100644 --- a/docs/history.md +++ b/docs/history.md @@ -2,6 +2,8 @@ ## 0.8.1 (Unreleased) +* ENH: Added georeferencing for Cartesian (x, y, z) and geographic (lat, lon, alt) coordinates in radar data processing ({issue}`243`) by [@syedhamidali](https://github.com/syedhamidali), ({pull}`244`) by [@syedhamidali](https://github.com/syedhamidali) + * ENH: Adding test to `open_datatree` function for all backends. Adding "scan_name" to nexradlevel2 datatree attributes ({pull}`238`) by [@aladinor](https://github.com/aladinor) * FIX: Improving performance of `open_nexradlevel2_datatree` function and adding tests for `sweep` parameter. ({issue}`239`) ({pull}`240`) by [@aladinor](https://github.com/aladinor) * FIX: Keeping attributes at each variable when using `open_nexradlevel2_datatree`. ({issue}`241`) ({pull}`242`) by [@aladinor](https://github.com/aladinor) diff --git a/docs/usage.md b/docs/usage.md index f34832cd..35b0121f 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -48,4 +48,5 @@ notebooks/angle_reindexing notebooks/Multi-Volume-Concatenation.ipynb notebooks/multiple-sweeps-into-volume-scan.ipynb notebooks/Transform +notebooks/georeferencing ``` diff --git a/examples/notebooks/georeferencing.ipynb b/examples/notebooks/georeferencing.ipynb new file mode 100644 index 00000000..dcb0b607 --- /dev/null +++ b/examples/notebooks/georeferencing.ipynb @@ -0,0 +1,396 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Georeferencing" + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "This notebook shows how to add Cartesian (x, y, z) and geographic (lat, lon, alt) coordinates to radar data using xradar.\n", + "\n", + "\n", + "- **DataTree** refers to [xarray.DataTree](https://docs.xarray.dev/en/stable/generated/xarray.DataTree.html#xarray.DataTree)\n", + "- **Dataset** refers to [xarray.Dataset](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.html#xarray.Dataset)\n", + "- **DataArray** refers to [xarray.DataArray](https://docs.xarray.dev/en/stable/generated/xarray.DataArray.html#xarray.DataArray)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "\n", + "import cmweather # noqa\n", + "import matplotlib.pyplot as plt\n", + "import matplotlib.ticker as ticker\n", + "from open_radar_data import DATASETS\n", + "\n", + "import xradar as xd\n", + "\n", + "warnings.simplefilter(\"ignore\")" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## Read Data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "file = DATASETS.fetch(\"sample_rainbow_5_59.vol\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "radar = xd.io.open_rainbow_datatree(file)" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "## Georeferencing (Cartesian)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "radar = radar.xradar.georeference()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "display(radar[\"sweep_0\"].ds)" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "- **Notice that the radar DataTree now includes x, y, and z coordinates.**\n", + "- **Next, let’s add lat, lon, and alt coordinates.**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "del radar" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## Georeferencing (Geographic)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "# Reading data again\n", + "radar = xd.io.open_rainbow_datatree(file)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "radar = radar.xradar.georeference(geo=True)\n", + "display(radar[\"sweep_0\"].ds)" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "- **When geo=True, lat, lon, and alt coordinates are added instead.**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "del radar" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "## Other Uses" + ] + }, + { + "cell_type": "markdown", + "id": "17", + "metadata": {}, + "source": [ + "- **You can also add these coordinates to a single DataArray or Dataset.**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "# Reading data again\n", + "radar = xd.io.open_rainbow_datatree(file)" + ] + }, + { + "cell_type": "markdown", + "id": "19", + "metadata": {}, + "source": [ + "## Assign both cartesian and geographic coords" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20", + "metadata": {}, + "outputs": [], + "source": [ + "radar = radar.xradar.georeference().xradar.georeference(geo=True)\n", + "display(radar[\"sweep_0\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "del radar" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "## Assign to Dataset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "# Reading data again\n", + "radar = xd.io.open_rainbow_datatree(file)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24", + "metadata": {}, + "outputs": [], + "source": [ + "ds = radar[\"sweep_0\"].to_dataset()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "display(ds)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26", + "metadata": {}, + "outputs": [], + "source": [ + "# Attach both types of coords\n", + "ds = ds.xradar.georeference().xradar.georeference(geo=True)\n", + "display(ds)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "ds = ds.drop(labels=[\"x\", \"y\", \"z\", \"lat\", \"lon\", \"alt\"])" + ] + }, + { + "cell_type": "markdown", + "id": "28", + "metadata": {}, + "source": [ + "## Assign to DataArray" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29", + "metadata": {}, + "outputs": [], + "source": [ + "reflectivity = ds[\"DBZH\"]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "reflectivity = reflectivity.xradar.georeference().xradar.georeference(geo=True)\n", + "display(reflectivity)" + ] + }, + { + "cell_type": "markdown", + "id": "31", + "metadata": {}, + "source": [ + "## Plot" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(1, 2, figsize=(10, 5.5))\n", + "\n", + "# Cartesian plot with x, y in km\n", + "reflectivity.plot(\n", + " x=\"x\",\n", + " y=\"y\",\n", + " add_colorbar=False,\n", + " vmin=-10,\n", + " vmax=60,\n", + " cmap=\"NWSRef\",\n", + " ax=ax[0],\n", + ")\n", + "ax[0].set_title(\"Reflectivity (Cartesian)\")\n", + "ax[0].xaxis.set_major_formatter(ticker.FuncFormatter(lambda val, pos: f\"{val/1e3:.0f}\"))\n", + "ax[0].yaxis.set_major_formatter(ticker.FuncFormatter(lambda val, pos: f\"{val/1e3:.0f}\"))\n", + "ax[0].set_xlabel(\"x (km)\")\n", + "ax[0].set_ylabel(\"y (km)\")\n", + "\n", + "# Geographic plot with lon, lat\n", + "pl = reflectivity.plot(\n", + " x=\"lon\",\n", + " y=\"lat\",\n", + " add_colorbar=False,\n", + " vmin=-10,\n", + " vmax=60,\n", + " cmap=\"NWSRef\",\n", + " ax=ax[1],\n", + ")\n", + "ax[1].set_title(\"Reflectivity (Geographic)\")\n", + "ax[1].set_xlabel(\"Longitude\")\n", + "ax[1].set_ylabel(\"Latitude\")\n", + "\n", + "# Add a centered colorbar below the plots\n", + "cbar = fig.colorbar(pl, ax=ax, orientation=\"horizontal\", fraction=0.05, pad=-0.25)\n", + "cbar.set_label(\"Reflectivity (dBZ)\")\n", + "\n", + "for ax in ax.flat:\n", + " ax.minorticks_on()\n", + " ax.grid(ls=\"--\", lw=0.5)\n", + "\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "33", + "metadata": {}, + "source": [ + "These transformations make radar data more versatile for analysis and visualization. You can apply these enhancements to entire DataTree structures or individual Dataset and DataArray objects, making xradar a flexible tool for radar data processing." + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/georeference/test_transforms.py b/tests/georeference/test_transforms.py index 3bd4a255..3069aeb7 100644 --- a/tests/georeference/test_transforms.py +++ b/tests/georeference/test_transforms.py @@ -7,7 +7,12 @@ from numpy.testing import assert_almost_equal import xradar -from xradar.georeference import antenna_to_cartesian, get_x_y_z +from xradar.georeference import ( + antenna_to_cartesian, + cartesian_to_geographic_aeqd, + get_lat_lon_alt, + get_x_y_z, +) def test_antenna_to_cartesian(): @@ -30,6 +35,26 @@ def test_antenna_to_cartesian(): assert_almost_equal(y[np.where(azimuths == 270.0)], 0) +def test_cartesian_to_geographic_aeqd(): + # Define test values + x = np.array([0, 1000, -1000]) + y = np.array([0, 500, -500]) + lon_0 = -97.59 + lat_0 = 36.49 + earth_radius = 6371000 # Earth's radius in meters + + # Convert Cartesian to geographic + lon, lat = cartesian_to_geographic_aeqd(x, y, lon_0, lat_0, earth_radius) + + # Check that the origin remains unchanged + assert_almost_equal(lon[0], lon_0) + assert_almost_equal(lat[0], lat_0) + + # Check that the coordinates are within a reasonable range + assert np.all(np.abs(lon) <= 180) + assert np.all(np.abs(lat) <= 90) + + def test_get_x_y_z(): # Create default xradar dataset ds = xradar.model.create_sweep_dataset() @@ -82,3 +107,25 @@ def test_get_x_y_z(): assert ds.crs_wkt.attrs[key] == pytest.approx(value) else: assert ds.crs_wkt.attrs[key] == value + + +def test_get_lat_lon_alt(): + # Create default xradar dataset + ds = xradar.model.create_sweep_dataset() + + # Apply lat, lon, and alt georeferencing + ds = get_lat_lon_alt(ds.swap_dims({"time": "azimuth"})) + + # Check that latitude, longitude, and altitude have been added + assert "lon" in ds.coords + assert "lat" in ds.coords + assert "alt" in ds.coords + + # # Check that the first range bin latitude and longitude are close to the radar location + origin = ds.isel(range=0).reset_coords().mean("azimuth") + assert_almost_equal(ds.latitude, origin.lat) + assert_almost_equal( + ds.longitude, + origin.lon, + ) + assert_almost_equal(ds.altitude, origin.alt - 0.8727675) diff --git a/tests/test_accessors.py b/tests/test_accessors.py index b31fa612..a8920744 100644 --- a/tests/test_accessors.py +++ b/tests/test_accessors.py @@ -14,40 +14,94 @@ def test_georeference_dataarray(): radar = xd.model.create_sweep_dataset() radar["sample_field"] = radar.azimuth + radar.range - geo = radar.sample_field.xradar.georeference() - assert_almost_equal(geo.x.values[:3, 0], np.array([0.436241, 1.3085901, 2.1805407])) + # geo=True + geo_cartesian = radar.sample_field.xradar.georeference(geo=False) assert_almost_equal( - geo.y.values[:3, 0], np.array([49.9882679, 49.973041, 49.9425919]) + geo_cartesian.x.values[:3, 0], np.array([0.436241, 1.3085901, 2.1805407]) ) assert_almost_equal( - geo.z.values[:3, 0], np.array([375.8727675, 375.8727675, 375.8727675]) + geo_cartesian.y.values[:3, 0], np.array([49.9882679, 49.973041, 49.9425919]) + ) + assert_almost_equal( + geo_cartesian.z.values[:3, 0], np.array([375.8727675, 375.8727675, 375.8727675]) + ) + + # geo=True + radar = xd.model.create_sweep_dataset().swap_dims({"time": "azimuth"}) + radar["sample_field"] = radar.azimuth + radar.range + + geo_geographic = radar.sample_field.xradar.georeference(geo=True) + assert_almost_equal( + geo_geographic.lon.values[:3, 0], np.array([8.7877327, 8.7877441, 8.7877554]) + ) + assert_almost_equal( + geo_geographic.lat.values[:3, 0], + np.array([46.1729908, 46.1729907, 46.17299042]), + ) + assert_almost_equal( + geo_geographic.alt.values[:3, 0], + np.array([375.87276751, 375.87276751, 375.87276751]), ) def test_georeference_dataset(): radar = xd.model.create_sweep_dataset() - geo = radar.xradar.georeference() - assert_almost_equal(geo.x.values[:3, 0], np.array([0.436241, 1.3085901, 2.1805407])) + geo_cartesian = radar.xradar.georeference(geo=False) + assert_almost_equal( + geo_cartesian.x.values[:3, 0], np.array([0.436241, 1.3085901, 2.1805407]) + ) + assert_almost_equal( + geo_cartesian.y.values[:3, 0], np.array([49.9882679, 49.973041, 49.9425919]) + ) + assert_almost_equal( + geo_cartesian.z.values[:3, 0], np.array([375.8727675, 375.8727675, 375.8727675]) + ) + + # geo=True + radar = xd.model.create_sweep_dataset().swap_dims({"time": "azimuth"}) + geo_geographic = radar.xradar.georeference(geo=True) assert_almost_equal( - geo.y.values[:3, 0], np.array([49.9882679, 49.973041, 49.9425919]) + geo_geographic.lon.values[:3, 0], np.array([8.7877327, 8.7877441, 8.7877554]) ) assert_almost_equal( - geo.z.values[:3, 0], np.array([375.8727675, 375.8727675, 375.8727675]) + geo_geographic.lat.values[:3, 0], + np.array([46.1729908, 46.1729907, 46.17299042]), + ) + assert_almost_equal( + geo_geographic.alt.values[:3, 0], + np.array([375.87276751, 375.87276751, 375.87276751]), ) def test_georeference_datatree(): radar = xd.model.create_sweep_dataset() tree = xr.DataTree.from_dict({"sweep_0": radar}) - geo = tree.xradar.georeference()["sweep_0"] + geo_cartesian = tree.xradar.georeference(geo=False)["sweep_0"] + assert_almost_equal( + geo_cartesian["x"].values[:3, 0], np.array([0.436241, 1.3085901, 2.1805407]) + ) + assert_almost_equal( + geo_cartesian["y"].values[:3, 0], np.array([49.9882679, 49.973041, 49.9425919]) + ) + assert_almost_equal( + geo_cartesian["z"].values[:3, 0], + np.array([375.8727675, 375.8727675, 375.8727675]), + ) + + # geo=True + radar = xd.model.create_sweep_dataset().swap_dims({"time": "azimuth"}) + tree = xr.DataTree.from_dict({"sweep_0": radar}) + geo_geographic = tree.xradar.georeference(geo=True)["sweep_0"] assert_almost_equal( - geo["x"].values[:3, 0], np.array([0.436241, 1.3085901, 2.1805407]) + geo_geographic.lon.values[:3, 0], np.array([8.7877327, 8.7877441, 8.7877554]) ) assert_almost_equal( - geo["y"].values[:3, 0], np.array([49.9882679, 49.973041, 49.9425919]) + geo_geographic.lat.values[:3, 0], + np.array([46.1729908, 46.1729907, 46.17299042]), ) assert_almost_equal( - geo["z"].values[:3, 0], np.array([375.8727675, 375.8727675, 375.8727675]) + geo_geographic.alt.values[:3, 0], + np.array([375.87276751, 375.87276751, 375.87276751]), ) diff --git a/xradar/accessors.py b/xradar/accessors.py index 2e08eee3..72b9287c 100644 --- a/xradar/accessors.py +++ b/xradar/accessors.py @@ -27,7 +27,14 @@ import xarray as xr -from .georeference import add_crs, add_crs_tree, get_crs, get_x_y_z, get_x_y_z_tree +from .georeference import ( + add_crs, + add_crs_tree, + get_crs, + get_lat_lon_alt, + get_x_y_z, + get_x_y_z_tree, +) from .transform import to_cfradial1, to_cfradial2 from .util import map_over_sweeps @@ -73,7 +80,7 @@ class XradarDataArrayAccessor(XradarAccessor): """Adds a number of xradar specific methods to xarray.DataArray objects.""" def georeference( - self, earth_radius=None, effective_radius_fraction=None + self, earth_radius=None, effective_radius_fraction=None, geo=False ) -> xr.DataArray: """ Parameters @@ -90,8 +97,13 @@ def georeference( """ radar = self.xarray_obj + if geo: + func = get_lat_lon_alt + else: + func = get_x_y_z + return radar.pipe( - get_x_y_z, + func, earth_radius=earth_radius, effective_radius_fraction=effective_radius_fraction, ) @@ -123,7 +135,7 @@ class XradarDataSetAccessor(XradarAccessor): """Adds a number of xradar specific methods to xarray.DataArray objects.""" def georeference( - self, earth_radius=None, effective_radius_fraction=None + self, earth_radius=None, effective_radius_fraction=None, geo=False ) -> xr.Dataset: """ Add georeference information to an xarray dataset @@ -140,8 +152,13 @@ def georeference( Dataset including x, y, and z as coordinates. """ radar = self.xarray_obj + if geo: + func = get_lat_lon_alt + else: + func = get_x_y_z + return radar.pipe( - get_x_y_z, + func, earth_radius=earth_radius, effective_radius_fraction=effective_radius_fraction, ) @@ -185,7 +202,7 @@ class XradarDataTreeAccessor(XradarAccessor): """Adds a number of xradar specific methods to xarray.DataTree objects.""" def georeference( - self, earth_radius=None, effective_radius_fraction=None + self, earth_radius=None, effective_radius_fraction=None, geo=False ) -> xr.DataTree: """ Add georeference information to an xradar datatree object @@ -202,10 +219,12 @@ def georeference( Datatree including x, y, and z as coordinates. """ radar = self.xarray_obj + return radar.pipe( get_x_y_z_tree, earth_radius=earth_radius, effective_radius_fraction=effective_radius_fraction, + geo=geo, ) def add_crs(self) -> xr.DataTree: diff --git a/xradar/georeference/transforms.py b/xradar/georeference/transforms.py index 20e3d1e2..d4fdf946 100644 --- a/xradar/georeference/transforms.py +++ b/xradar/georeference/transforms.py @@ -13,11 +13,19 @@ {} """ -__all__ = ["antenna_to_cartesian", "get_x_y_z", "get_x_y_z_tree"] +__all__ = [ + "antenna_to_cartesian", + "cartesian_to_geographic_aeqd", + "get_x_y_z", + "get_lat_lon_alt", + "get_x_y_z_tree", +] __doc__ = __doc__.format("\n ".join(__all__)) +import warnings + import numpy as np import xarray as xr @@ -101,6 +109,65 @@ def antenna_to_cartesian( return x, y, z +def cartesian_to_geographic_aeqd(x, y, lon_0, lat_0, earth_radius): + """ + Transform Cartesian coordinates (x, y) to geographic coordinates (latitude, longitude) + using the Azimuthal Equidistant (AEQD) map projection. + + Parameters + ---------- + x, y : array-like + Cartesian coordinates in the same units as the Earth's radius, typically meters. + lon_0, lat_0 : float + Longitude and latitude, in degrees, of the center of the projection. + earth_radius : float + Radius of the Earth in meters. + + Returns + ------- + lon, lat : array + Longitude and latitude of Cartesian coordinates in degrees. + + Notes + ----- + The calculations follow the AEQD projection equations, where the Earth's radius is used + to define the distance metric. + """ + # Ensure x and y are at least 1D arrays + x = np.atleast_1d(np.asarray(x)) + y = np.atleast_1d(np.asarray(y)) + + # Convert reference latitude and longitude to radians + lat_0_rad = np.deg2rad(lat_0) + lon_0_rad = np.deg2rad(lon_0) + + # Calculate distance (rho) and angular distance (c) + rho = np.sqrt(x**2 + y**2) + c = rho / earth_radius + + # Suppress warnings for potential division by zero + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + lat_rad = np.arcsin( + np.cos(c) * np.sin(lat_0_rad) + (y * np.sin(c) * np.cos(lat_0_rad) / rho) + ) + + # Convert latitude to degrees and handle edge cases where rho == 0 + lat_deg = np.rad2deg(lat_rad) + lat_deg[rho == 0] = lat_0 + + # Calculate longitude in radians + x1 = x * np.sin(c) + x2 = rho * np.cos(lat_0_rad) * np.cos(c) - y * np.sin(lat_0_rad) * np.sin(c) + lon_rad = lon_0_rad + np.arctan2(x1, x2) + + # Convert longitude to degrees and normalize to [-180, 180] + lon_deg = np.rad2deg(lon_rad) + lon_deg = (lon_deg + 180) % 360 - 180 # Normalize longitude + + return lon_deg, lat_deg + + def get_x_y_z(ds, earth_radius=None, effective_radius_fraction=None): """ Return Cartesian coordinates from antenna coordinates. @@ -179,30 +246,127 @@ def get_x_y_z(ds, earth_radius=None, effective_radius_fraction=None): return ds -def get_x_y_z_tree(radar, earth_radius=None, effective_radius_fraction=None): +# def get_x_y_z_tree(radar, earth_radius=None, effective_radius_fraction=None): +# """ +# Applies the georeferencing to a xradar datatree + +# Parameters +# ---------- +# radar: xarray.DataTree +# Xradar datatree object with radar information. +# earth_radius: float +# Radius of the earth. Defaults to a latitude-dependent radius derived from +# WGS84 ellipsoid. +# effective_radius_fraction: float +# Fraction of earth to use for the effective radius (default is 4/3). + +# Returns +# ------- +# radar: xarray.DataTree +# Datatree with sweep datasets including georeferenced coordinates +# """ +# for key in list(radar.children): +# if "sweep" in key: +# radar[key].ds = get_x_y_z( +# radar[key].to_dataset(), +# earth_radius=earth_radius, +# effective_radius_fraction=effective_radius_fraction, +# ) +# return radar + + +def get_lat_lon_alt(ds, earth_radius=None, effective_radius_fraction=None): """ - Applies the georeferencing to a xradar datatree + Add geographic coordinates (lat, lon, alt) to the dataset. Parameters ---------- - radar: xarray.DataTree + ds : xarray.Dataset + Xarray dataset containing range, azimuth, and elevation. + earth_radius : float + Radius of the earth. Defaults to a latitude-dependent radius derived from + WGS84 ellipsoid. + effective_radius_fraction : float + Fraction of earth to use for the effective radius (default is 4/3). + + Returns + ------- + ds : xarray.Dataset + Dataset including lat, lon, and alt as coordinates. + """ + if earth_radius is None: + crs = get_crs(ds) + earth_radius = get_earth_radius(crs, ds.latitude.values) + ds = ds.pipe(add_crs, crs) + + # Calculate Cartesian coordinates first + x, y, z = antenna_to_cartesian( + ds.range, + ds.azimuth, + ds.elevation, + earth_radius=earth_radius, + effective_radius_fraction=effective_radius_fraction, + site_altitude=ds.altitude.values, + ) + + # Convert Cartesian coordinates to geographic coordinates + lon, lat = cartesian_to_geographic_aeqd( + x, y, ds.longitude.values, ds.latitude.values, earth_radius=earth_radius + ) + # alt = z - ds.altitude.values + alt = z + + # Add coordinates to the dataset + ds["lat"] = xr.DataArray(lat, dims=["azimuth", "range"]) + ds["lat"].attrs = {"standard_name": "latitude", "units": "degrees_north"} + + ds["lon"] = xr.DataArray(lon, dims=["azimuth", "range"]) + ds["lon"].attrs = {"standard_name": "longitude", "units": "degrees_east"} + + ds["alt"] = xr.DataArray(alt, dims=["azimuth", "range"]) + ds["alt"].attrs = {"standard_name": "altitude", "units": "meters"} + + # Make sure the coordinates are set properly + if isinstance(ds, xr.Dataset): + ds = ds.set_coords(["lat", "lon", "alt"]) + + return ds + + +def get_x_y_z_tree(radar, earth_radius=None, effective_radius_fraction=None, geo=False): + """ + Applies the georeferencing to a xradar datatree. + + Parameters + ---------- + radar : xarray.DataTree Xradar datatree object with radar information. - earth_radius: float + earth_radius : float Radius of the earth. Defaults to a latitude-dependent radius derived from WGS84 ellipsoid. - effective_radius_fraction: float + effective_radius_fraction : float Fraction of earth to use for the effective radius (default is 4/3). + geo : bool, optional + If True, adds geographic coordinates (lat, lon, alt) instead of Cartesian + coordinates (x, y, z). Default is False. Returns ------- - radar: xarray.DataTree - Datatree with sweep datasets including georeferenced coordinates + radar : xarray.DataTree + Datatree with sweep datasets including georeferenced coordinates. """ for key in list(radar.children): if "sweep" in key: - radar[key].ds = get_x_y_z( - radar[key].to_dataset(), - earth_radius=earth_radius, - effective_radius_fraction=effective_radius_fraction, - ) + if geo: + radar[key].ds = get_lat_lon_alt( + radar[key].to_dataset(), + earth_radius=earth_radius, + effective_radius_fraction=effective_radius_fraction, + ) + else: + radar[key].ds = get_x_y_z( + radar[key].to_dataset(), + earth_radius=earth_radius, + effective_radius_fraction=effective_radius_fraction, + ) return radar