Source code for decode.convert
__all__ = ["coord_units", "frame", "units"]
# standard library
from typing import Optional, Sequence, TypeVar, Union
# dependencies
import numpy as np
import xarray as xr
from astropy.units import Equivalency, Quantity, Unit
# type hints
T = TypeVar("T")
Multiple = Union[Sequence[T], T]
UnitLike = Union[xr.DataArray, Unit, str]
[docs]
def coord_units(
da: xr.DataArray,
coord_names: Multiple[str],
new_units: UnitLike,
/,
equivalencies: Optional[Equivalency] = None,
) -> xr.DataArray:
"""Convert the units of coordinate(s) of a DataArray.
Args:
da: Input DataArray.
coord_names: Name(s) of the coordinate(s) to be converted.
new_units: Units to be converted from the current ones.
A DataArray that has units attribute is also accepted.
equivalencies: Optional Astropy equivalencies.
Returns:
DataArray with the units of the coordinate converted.
"""
# deepcopy except for data
da = da.copy(data=da.data)
if isinstance(coord_names, str):
coord_names = [coord_names]
for coord_name in coord_names:
coord = da.coords[coord_name]
new_coord = units(coord, new_units, equivalencies)
da = da.assign_coords({coord_name: new_coord})
return da
[docs]
def frame(da: xr.DataArray, new_frame: str, /) -> xr.DataArray:
"""Convert the skycoord frame of a DataArray.
Args:
da: Input DataArray.
new_frame: Frame to be converted from the current one.
Returns:
DataArray with the skycoord frame converted.
"""
# deepcopy except for data
da = da.copy(data=da.data)
if not new_frame == "relative":
raise ValueError("Relative is only available.")
lon = da.coords["lon"]
lat = da.coords["lat"]
lon_origin = da.coords["lon_origin"]
lat_origin = da.coords["lat_origin"]
# do not change the order below!
lon -= units(lon_origin, lon)
lon *= np.cos(units(lat, "rad"))
lat -= units(lat_origin, lat)
lon_origin *= 0.0
lat_origin *= 0.0
new_frame = da.frame.copy(data=new_frame)
return da.assign_coords(frame=new_frame)
[docs]
def units(
da: xr.DataArray,
new_units: UnitLike,
/,
equivalencies: Optional[Equivalency] = None,
) -> xr.DataArray:
"""Convert the units of a DataArray.
Args:
da: Input DataArray.
new_units: Units to be converted from the current ones.
A DataArray that has units attribute is also accepted.
equivalencies: Optional Astropy equivalencies.
Returns:
DataArray with the units converted.
"""
if (units := da.attrs.get("units")) is None:
raise ValueError("Units must exist in DataArray attrs.")
if isinstance(new_units, xr.DataArray):
new_units = new_units.attrs["units"]
new_data = Quantity(da, units).to(new_units, equivalencies)
return da.copy(data=new_data).assign_attrs(units=new_units)