diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 71474120..bcf0fe82 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -78,7 +78,7 @@ jobs: - name: install dependencies run: | - python -m pip install -r ci/requirements.txt + python -m pip install ".[test]" - name: install astropy-xarray run: python -m pip install --no-deps . diff --git a/.gitignore b/.gitignore index 0cedeb7e..1ce64b54 100644 --- a/.gitignore +++ b/.gitignore @@ -104,7 +104,7 @@ celerybeat.pid # Environments .env -.venv +.venv*/ env/ venv/ ENV/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3d8f4987..25a0be3a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -39,7 +39,7 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.12.3 hooks: - - id: ruff + - id: ruff-check args: [--fix] - repo: https://github.com/kynan/nbstripout rev: 0.8.1 diff --git a/astropy_xarray/__init__.py b/astropy_xarray/__init__.py index 8f17bca7..13e6702c 100644 --- a/astropy_xarray/__init__.py +++ b/astropy_xarray/__init__.py @@ -17,6 +17,7 @@ import astropy.units +import astropy_xarray.time_compat # noqa: F401 from astropy_xarray import accessors, formatting, testing # noqa: F401 from astropy_xarray.index import AstropyIndex diff --git a/astropy_xarray/accessors.py b/astropy_xarray/accessors.py index 58ae4619..7b58ff18 100644 --- a/astropy_xarray/accessors.py +++ b/astropy_xarray/accessors.py @@ -18,11 +18,16 @@ import astropy import astropy.units -from xarray import register_dataarray_accessor, register_dataset_accessor +from xarray import ( + register_dataarray_accessor, + register_dataset_accessor, + register_datatree_accessor, +) from xarray.core.dtypes import NA from astropy_xarray import conversion -from astropy_xarray.conversion import no_unit_values +from astropy_xarray.conversion import AstropyUnitType, no_unit_values +from astropy_xarray.coordinates.sky_coord import dataset_to_skycoord from astropy_xarray.errors import format_error_message # sentinel to fallback to attribute unit, then value container type @@ -64,7 +69,7 @@ def units_to_str_or_none(mapping, unit_format): formatter = str if not unit_format else lambda v: unit_format.format(v) return { - key: formatter(value) if isinstance(value, astropy.units.UnitBase) else value + key: formatter(value) if isinstance(value, AstropyUnitType) else value for key, value in mapping.items() } @@ -96,6 +101,8 @@ def _decide_unit(unit, unit_attribute): elif unit is _default: if unit_attribute in no_unit_values: return unit_attribute + if isinstance(unit_attribute, dict): + return unit_attribute if isinstance(unit_attribute, astropy.units.UnitBase): unit = unit_attribute else: @@ -1718,3 +1725,47 @@ def interpolate_na( ) return conversion.attach_units(interpolated, units) + + def to_sky_coord(self): + """TODO""" + return dataset_to_skycoord(self.ds) + + +@register_datatree_accessor("astropy") +class AstropyDataTreeAccessor: + """ + Access methods for DataTree with units using Astropy. + + Methods and attributes can be accessed through the `.astropy` attribute. + """ + + def __init__(self, dt): + self.dt = dt + + def quantify(self, units=_default, **unit_kwargs): + from xarray import DataTree + + def breadth_first_quantify(dt: DataTree): + quantified_ds = ( + dt.dataset.astropy.quantify(units=_default, **unit_kwargs) + if dt.dataset is not None + else None + ) + children = {k: breadth_first_quantify(v) for k, v in dt.children.items()} + return DataTree(dataset=quantified_ds, children=children) + + return breadth_first_quantify(self.dt) + + def dequantify(self, format=None): + from xarray import DataTree + + def breadth_first_dequantify(dt: DataTree): + dequantified_ds = ( + dt.dataset.astropy.dequantify(format=format) + if dt.dataset is not None + else None + ) + children = {k: breadth_first_dequantify(v) for k, v in dt.children.items()} + return DataTree(dataset=dequantified_ds, children=children) + + return breadth_first_dequantify(self.dt) diff --git a/astropy_xarray/conversion.py b/astropy_xarray/conversion.py index de01d391..79538959 100644 --- a/astropy_xarray/conversion.py +++ b/astropy_xarray/conversion.py @@ -17,6 +17,8 @@ import re import astropy +import astropy.coordinates +import astropy.time import astropy.units from xarray import Coordinates, DataArray, Dataset, IndexVariable, Variable @@ -34,11 +36,15 @@ datetime_units_re = re.compile(rf"{time_units_re} since {datetime_re}") +AstropyType = astropy.units.Quantity | astropy.time.TimeBase +AstropyUnitType = astropy.units.UnitBase | astropy.units.FunctionUnitBase + + def is_datetime_unit(unit): return isinstance(unit, str) and datetime_units_re.match(unit) is not None -def array_attach_unit(data, unit) -> astropy.units.Quantity: +def array_attach_unit(data, unit) -> AstropyType: """attach a unit to the data Parameters @@ -55,7 +61,7 @@ def array_attach_unit(data, unit) -> astropy.units.Quantity: if unit in no_unit_values: return data - if not isinstance(unit, astropy.units.UnitBase): + if not isinstance(unit, AstropyUnitType | dict): raise ValueError(f"cannot use {unit!r} as a unit") if isinstance(data, astropy.units.Quantity): @@ -67,7 +73,35 @@ def array_attach_unit(data, unit) -> astropy.units.Quantity: f"already has units {data.unit}" ) - return astropy.units.Quantity(data, unit) + if isinstance(unit, dict): + match unit["class"].lower(): + case "time": + return astropy.time.Time( + data, + format=unit["format"], + scale=unit["scale"], + precision=unit["precision"], + ) + case "timedelta": + return astropy.time.TimeDelta( + data, + format=unit["format"], + scale=unit["scale"], + precision=unit["precision"], + ) + case "angle": + return astropy.coordinates.Angle(data, unit=unit["unit"]) + case "longitude": + return astropy.coordinates.Longitude(data, unit=unit["unit"]) + case "latitude": + return astropy.coordinates.Latitude(data, unit=unit["unit"]) + case "distance": + return astropy.coordinates.Distance(data, unit=unit["unit"]) + + if isinstance(unit, astropy.units.LogUnit): + return astropy.units.LogQuantity(data, unit) + else: + return astropy.units.Quantity(data, unit) def array_convert_unit(data, unit, equivalencies) -> astropy.units.Quantity: @@ -111,13 +145,32 @@ def array_convert_unit(data, unit, equivalencies) -> astropy.units.Quantity: return data -def array_extract_unit(data): +def array_extract_unit(data) -> dict | astropy.units.Unit | astropy.units.LogUnit: """extract the unit of an array If ``data`` is not a quantity, the units are ``None`` """ try: - return data.unit + if isinstance( + data, + ( + astropy.coordinates.Longitude, + astropy.coordinates.Latitude, + astropy.coordinates.Distance, + ), + ): + return {"class": data.__class__.__name__.lower(), "unit": str(data.unit)} + elif isinstance(data, astropy.time.TimeBase): + return { + "class": data.__class__.__name__.lower(), + "format": data.format, + "scale": data.scale, + "precision": data.precision, + } + elif isinstance(data, astropy.units.Quantity): + return data.unit + else: + return None except AttributeError: return None @@ -125,7 +178,10 @@ def array_extract_unit(data): def array_strip_unit(data): """strip the unit of a quantity""" try: - return data.value + if isinstance(data, (astropy.units.Quantity, astropy.time.TimeBase)): + return data.value + else: + return data except AttributeError: return data @@ -405,7 +461,7 @@ def extract_unit_attributes(obj, attr="units"): def strip_units_variable(var): - if not isinstance(var.data, astropy.units.Quantity): + if not isinstance(var.data, AstropyType): return var data = array_strip_unit(var.data) diff --git a/astropy_xarray/coordinates/__init__.py b/astropy_xarray/coordinates/__init__.py new file mode 100644 index 00000000..12712059 --- /dev/null +++ b/astropy_xarray/coordinates/__init__.py @@ -0,0 +1,12 @@ +from astropy_xarray.coordinates.frame import load_frame, load_representation +from astropy_xarray.coordinates.sky_coord import ( + dataset_to_skycoord, + skycoord_to_dataset, +) + +__all__ = [ + "dataset_to_skycoord", + "skycoord_to_dataset", + "load_frame", + "load_representation", +] diff --git a/astropy_xarray/coordinates/core.py b/astropy_xarray/coordinates/core.py new file mode 100644 index 00000000..de34db01 --- /dev/null +++ b/astropy_xarray/coordinates/core.py @@ -0,0 +1,26 @@ +from typing import TypeVar + +import astropy.units as u +from astropy.time import Time + + +# Time +def dump_time(time: Time): + return { + "val": time.value, + "format": time.format, + "precision": time.precision, + "scale": time.scale, + } + + +# Quantity +def dump_quantity(q: u.Quantity): + return {"value": float(q.value), "unit": str(q.unit)} + + +_T = TypeVar("_T", bound=u.Quantity | Time) + + +def load_optional_object(cls: type[_T], kwargs: dict | None) -> _T | None: + return cls(**kwargs) if kwargs is not None else None diff --git a/astropy_xarray/coordinates/frame.py b/astropy_xarray/coordinates/frame.py new file mode 100644 index 00000000..54f3436e --- /dev/null +++ b/astropy_xarray/coordinates/frame.py @@ -0,0 +1,325 @@ +from types import NoneType + +import astropy.units as u +import numpy as np +from astropy.coordinates import ( + BaseCoordinateFrame, + BaseRepresentation, + CartesianDifferential, + CartesianRepresentation, + CylindricalRepresentation, + EarthLocation, + SphericalRepresentation, + frame_transform_graph, + representation, +) +from astropy.coordinates.builtin_frames import ( + CIRS, + FK4, + FK5, + GCRS, + HCRS, + ICRS, + ITRS, + AltAz, + CustomBarycentricEcliptic, + Galactic, + GalacticLSR, + Galactocentric, + GeocentricMeanEcliptic, + GeocentricTrueEcliptic, + HADec, + HeliocentricEclipticIAU76, + HeliocentricMeanEcliptic, + HeliocentricTrueEcliptic, + PrecessedGeocentric, + Supergalactic, +) +from astropy.time import Time + +from astropy_xarray.coordinates.core import ( + dump_quantity, + dump_time, + load_optional_object, +) +from astropy_xarray.coordinates.representation import ( + dump_cartesian_differential, + dump_cartesian_representation, + dump_cylindrical_representation, + dump_earth_location, + dump_spherical_representation, + load_cartesian_differential, + load_cartesian_representation, + load_optional_earthlocation, +) + + +def dump_frame(frame: BaseCoordinateFrame, with_data: bool = False) -> dict: + ser = { + "name": frame.name, + "representation_type": frame.representation_type.get_name(), + "differential_type": frame.differential_type.get_name(), + } + if frame.has_data: + ser["data"] = {"representation_type": frame.data.get_name()} + if "s" in frame.data.differentials: + ser["data"]["differential_type"] = frame.data.differentials["s"].get_name() + if with_data: + for component in frame.data.components: + quantity: u.Quantity = getattr(frame.data, component) + ser["data"][component] = { + "value": float(quantity.value), + "unit": str(quantity.unit), + } + if "s" in frame.data.differentials: + for component in frame.data.differentials["s"].components: + quantity: u.Quantity = getattr( + frame.data.differentials["s"], component + ) + ser["data"][component] = { + "value": float(quantity.value), + "unit": str(quantity.unit), + } + + for attribute_name in frame.frame_attributes.keys(): + attr = ser[attribute_name] = getattr(frame, attribute_name) + if not isinstance(attr, NoneType): + match attr: + case BaseCoordinateFrame(): + ser_attr = dump_frame(attr, with_data=True) + case Time(): + ser_attr = dump_time(attr) + case EarthLocation(): + ser_attr = dump_earth_location(attr) + case SphericalRepresentation(): + ser_attr = dump_spherical_representation(attr) + case CartesianRepresentation(): + ser_attr = dump_cartesian_representation(attr) + case CylindricalRepresentation(): + ser_attr = dump_cylindrical_representation(attr) + case CartesianDifferential(): + ser_attr = dump_cartesian_differential(attr) + case u.Quantity(): + ser_attr = dump_quantity(attr) + case _: + raise NotImplementedError("unsupported frame member", attr) + ser[attribute_name] = ser_attr + return ser + + +def load_frame(frame_dict: dict, with_data: bool = False) -> BaseCoordinateFrame: + representation_type = frame_dict.get("data").get("representation_type") + differential_type = frame_dict.get("data").get("differential_type") + kwargs = dict( + representation_type=representation_type, differential_type=differential_type + ) + match frame_dict["name"]: + case "icrs": + frame = ICRS(**kwargs) + case "fk5": + frame = FK5( + equinox=load_optional_object(Time, frame_dict["equinox"]), **kwargs + ) + case "fk4": + frame = FK4( + equinox=load_optional_object(Time, frame_dict["equinox"]), + obstime=load_optional_object(Time, frame_dict["obstime"]), + **kwargs, + ) + case "galactic": + frame = Galactic(**kwargs) + case "galacticlsr": + frame = GalacticLSR(**kwargs) + case "galactocentric": + frame = Galactocentric( + galcen_distance=u.Quantity(**frame_dict["galcen_distance"]), + galcen_coord=load_frame(frame_dict["galcen_coord"], with_data=True), + galcen_v_sun=load_cartesian_differential(frame_dict["galcen_v_sun"]), + z_sun=u.Quantity(**frame_dict["z_sun"]), + roll=u.Quantity(**frame_dict["roll"]), + **kwargs, + ) + case "gcrs": + frame = GCRS( + obstime=load_optional_object(Time, frame_dict["obstime"]), + obsgeoloc=load_cartesian_representation(frame_dict["obsgeoloc"]), + obsgeovel=load_cartesian_representation(frame_dict["obsgeovel"]), + **kwargs, + ) + case "cirs": + frame = CIRS( + obstime=load_optional_object(Time, frame_dict["obstime"]), + location=load_optional_earthlocation(frame_dict["location"]), + **kwargs, + ) + case "itrs": + frame = ITRS( + obstime=load_optional_object(Time, frame_dict["obstime"]), + location=load_optional_earthlocation(frame_dict["location"]), + **kwargs, + ) + case "hcrs": + frame = HCRS( + obstime=load_optional_object(Time, frame_dict["obstime"]), **kwargs + ) + case "itrs": + frame = ITRS( + location=load_optional_earthlocation(frame_dict["location"]), + **kwargs, + ) + case "altaz": + frame = AltAz( + obstime=load_optional_object(Time, frame_dict["obstime"]), + location=load_optional_earthlocation(frame_dict["location"]), + pressure=load_optional_object(u.Quantity, frame_dict["pressure"]), + temperature=load_optional_object(u.Quantity, frame_dict["temperature"]), + relative_humidity=load_optional_object( + u.Quantity, frame_dict["relative_humidity"] + ), + obswl=load_optional_object(u.Quantity, frame_dict["obswl"]), + **kwargs, + ) + case "hadec": + frame = HADec( + obstime=load_optional_object(Time, frame_dict["obstime"]), + location=load_optional_earthlocation(frame_dict["location"]), + pressure=load_optional_object(u.Quantity, frame_dict["pressure"]), + temperature=load_optional_object(u.Quantity, frame_dict["temperature"]), + relative_humidity=load_optional_object( + u.Quantity, frame_dict["relative_humidity"] + ), + obswl=load_optional_object(u.Quantity, frame_dict["obswl"]), + **kwargs, + ) + case "supergalactic": + frame = Supergalactic(**kwargs) + case "precessedgeocentric": + frame = PrecessedGeocentric( + equinox=load_optional_object(Time, frame_dict["equinox"]), + obstime=load_optional_object(Time, frame_dict["obstime"]), + obsgeoloc=load_cartesian_representation(frame_dict["obsgeoloc"]), + obsgeovel=load_cartesian_representation(frame_dict["obsgeovel"]), + **kwargs, + ) + case "geocentricmeanecliptic": + frame = GeocentricMeanEcliptic( + equinox=load_optional_object(Time, frame_dict["equinox"]), + obstime=load_optional_object(Time, frame_dict["obstime"]), + **kwargs, + ) + case "geocentrictrueecliptic": + frame = GeocentricTrueEcliptic( + equinox=load_optional_object(Time, frame_dict["equinox"]), + obstime=load_optional_object(Time, frame_dict["obstime"]), + **kwargs, + ) + case "heliocentricmeanecliptic": + frame = HeliocentricMeanEcliptic( + equinox=load_optional_object(Time, frame_dict["equinox"]), + obstime=load_optional_object(Time, frame_dict["obstime"]), + **kwargs, + ) + case "heliocentrictrueecliptic": + frame = HeliocentricTrueEcliptic( + equinox=load_optional_object(Time, frame_dict["equinox"]), + obstime=load_optional_object(Time, frame_dict["obstime"]), + **kwargs, + ) + case "heliocentriceclipticiau76": + frame = HeliocentricEclipticIAU76( + obstime=load_optional_object(Time, frame_dict["obstime"]), **kwargs + ) + case "custombarycentricecliptic": + frame = CustomBarycentricEcliptic( + obliquity=load_optional_object(u.Quantity, frame_dict["obliquity"]), + **kwargs, + ) + case _: + raise NotImplementedError(frame_dict["name"]) + if with_data: + frame = frame.realize_frame( + representation_type=frame_dict.get("representation_type"), + differential_type=frame_dict.get("differential_type"), + data=load_representation( + representation_type, + differential_type, + None, + { + k: u.Quantity(**v) + for k, v in frame_dict.get("data").items() + if isinstance(v, dict) + }, + ), + ) + return frame + + +def load_representation( + representation_type: str, + differential_type: str | None, + frame_name: str | None, + data: dict[str, np.ndarray], +) -> BaseRepresentation: + RepresentationClass = representation.REPRESENTATION_CLASSES.get(representation_type) + DifferentialClass = representation.DIFFERENTIAL_CLASSES.get(differential_type) + + if frame_name is None: + # using data component names + differentials = ( + DifferentialClass( + **{key: data[key] for key in DifferentialClass.attr_classes.keys()}, + copy=False, + ) + if DifferentialClass is not None + else None + ) + + return RepresentationClass( + **{key: data[key] for key in RepresentationClass.attr_classes.keys()}, + differentials=differentials, + copy=False, + ) + else: + # using frame component names + frame_type: BaseCoordinateFrame = frame_transform_graph.lookup_name(frame_name) + diff_to_data = ( + dict( + zip( + tuple(DifferentialClass.attr_classes), + frame_type._get_representation_info()[DifferentialClass]["names"], + ) + ) + if DifferentialClass is not None + else {} + ) + rep_to_data = ( + dict( + zip( + tuple(RepresentationClass.attr_classes), + frame_type._get_representation_info()[RepresentationClass]["names"], + ) + ) + if RepresentationClass is not None + else {} + ) + + differentials = ( + DifferentialClass( + **{ + key: data[diff_to_data[key]] + for key in DifferentialClass.attr_classes.keys() + }, + copy=False, + ) + if DifferentialClass is not None + else None + ) + + return RepresentationClass( + **{ + key: data[rep_to_data[key]] + for key in RepresentationClass.attr_classes.keys() + }, + differentials=differentials, + copy=False, + ) diff --git a/astropy_xarray/coordinates/representation.py b/astropy_xarray/coordinates/representation.py new file mode 100644 index 00000000..7738a32e --- /dev/null +++ b/astropy_xarray/coordinates/representation.py @@ -0,0 +1,81 @@ +import astropy.units as u +from astropy.coordinates import ( + CartesianDifferential, + CartesianRepresentation, + CylindricalRepresentation, + EarthLocation, + SphericalRepresentation, +) + +from astropy_xarray.coordinates.core import dump_quantity + + +# Representation +def dump_spherical_representation(rep: SphericalRepresentation): + return { + "lon": dump_quantity(rep.lon), + "lat": dump_quantity(rep.lat), + "distance": dump_quantity(rep.distance), + } + + +def dump_cartesian_representation(rep: CartesianRepresentation): + return { + "x": dump_quantity(rep.x), + "y": dump_quantity(rep.y), + "z": dump_quantity(rep.z), + } + + +def load_cartesian_representation(kwargs: dict): + return CartesianRepresentation( + x=u.Quantity(**kwargs["x"]), + y=u.Quantity(**kwargs["y"]), + z=u.Quantity(**kwargs["z"]), + ) + + +def dump_cylindrical_representation(rep: CylindricalRepresentation): + return { + "rho": dump_quantity(rep.rho), + "phi": dump_quantity(rep.phi), + "z": dump_quantity(rep.z), + } + + +# Earth Location +def dump_earth_location(el: EarthLocation): + return { + "x": dump_quantity(el.x), + "y": dump_quantity(el.y), + "z": dump_quantity(el.z), + } + + +def load_optional_earthlocation(kwargs: dict | None) -> EarthLocation | None: + return ( + EarthLocation( + x=u.Quantity(**kwargs["x"]), + y=u.Quantity(**kwargs["y"]), + z=u.Quantity(**kwargs["z"]), + ) + if kwargs is not None + else None + ) + + +# Cartesian Differential +def dump_cartesian_differential(cd: CartesianDifferential): + return { + "d_x": dump_quantity(cd.d_x), + "d_y": dump_quantity(cd.d_y), + "d_z": dump_quantity(cd.d_z), + } + + +def load_cartesian_differential(kwargs: dict): + return CartesianDifferential( + d_x=u.Quantity(**kwargs["d_x"]), + d_y=u.Quantity(**kwargs["d_y"]), + d_z=u.Quantity(**kwargs["d_z"]), + ) diff --git a/astropy_xarray/coordinates/sky_coord.py b/astropy_xarray/coordinates/sky_coord.py new file mode 100644 index 00000000..6a1c2472 --- /dev/null +++ b/astropy_xarray/coordinates/sky_coord.py @@ -0,0 +1,123 @@ +import itertools +from collections.abc import Generator + +import astropy.units as u +import numpy as np +import xarray as xr +from astropy.coordinates import SkyCoord + +from astropy_xarray.coordinates.frame import dump_frame, load_frame, load_representation + + +def _skycoord_representation_component_names( + skycoord: SkyCoord, +) -> tuple[str, ...]: + frame = skycoord.frame.replicate_without_data(copy=True) + frame.representation_type = type(skycoord.frame.data) + if skycoord.frame.data.differentials: + frame.differential_type = type(skycoord.frame.data.differentials.get("s")) + + return tuple(name for name in frame.get_representation_component_names()) + + +def _skycoord_differential_component_names( + skycoord: SkyCoord, +) -> tuple[str, ...]: + if skycoord.data.differentials.get("s") is None: + return () + + frame = skycoord.frame.replicate_without_data(copy=True) + frame.representation_type = type(skycoord.frame.data) + if skycoord.frame.data.differentials: + frame.differential_type = type(skycoord.frame.data.differentials.get("s")) + + return tuple(name for name in frame.get_representation_component_names("s")) + + +def _skycoord_component_names( + skycoord: SkyCoord, +) -> Generator[str, None, None]: + yield from itertools.chain( + _skycoord_representation_component_names(skycoord), + _skycoord_differential_component_names(skycoord), + ) + # use default frame components + + +def _skycoord_components(skycoord: SkyCoord) -> dict[str, u.Quantity]: + frame = skycoord.frame.replicate_without_data(copy=True) + frame.representation_type = type(skycoord.frame.data) + if skycoord.frame.data.differentials: + frame.differential_type = type(skycoord.frame.data.differentials.get("s")) + data = skycoord.frame.data + + # need to access via data components to avoid rounding + frame_to_rep = ( + frame.get_representation_component_names() + | frame.get_representation_component_names("s") + ) + return { + name: ( + getattr(data, frame_to_rep[name]) + if hasattr(data, frame_to_rep[name]) + else getattr(data.differentials["s"], frame_to_rep[name]) + ) + for name in _skycoord_component_names(skycoord) + } + + +def _skycoord_to_dataarrays( + skycoord: SkyCoord, + coords: list[tuple[str, list]] | None = None, +) -> dict[str, xr.DataArray]: + return { + name: xr.DataArray( + data=component, coords=dict(coords) if coords is not None else None + ) + for name, component in _skycoord_components(skycoord).items() + } + + +def skycoord_to_dataset( + skycoord: SkyCoord, + coords: list[tuple[str, np.ndarray]] | None = None, +) -> xr.Dataset: + """Convert a SkyCoord object to an xarray Dataset. + + Args: + skycoord: SkyCoord object. + coords: coordinates to assign to multidimensional skycoords. Defaults to None. + + Returns: + quantified dataset. + """ + return xr.Dataset( + coords=dict(coords) if coords is not None else None, + data_vars=_skycoord_to_dataarrays(skycoord, coords), + attrs=dict( + frame=dump_frame(skycoord.frame), + ), + ) + + +def dataset_to_skycoord(ds: xr.Dataset) -> SkyCoord: + """Convert a SkyCoord-based Dataset with metadata attributes to a SkyCoord. + + Args: + ds: skycoord dataset. + use_frame_names: True if frame component names are used a data_vars. Defaults to True. + + Returns: + SkyCoord object. + """ + dsq = ds.astropy.quantify() + frame = load_frame(ds.attrs["frame"]) + frame._data = load_representation( + ds.attrs["frame"]["data"]["representation_type"], + ds.attrs["frame"]["data"].get("differential_type"), + ds.attrs["frame"]["name"], + {k: v.data for k, v in dsq.data_vars.items()}, + ) + frame.representation_type = ds.attrs["frame"]["representation_type"] + frame.differential_type = ds.attrs["frame"]["differential_type"] + return SkyCoord(frame) diff --git a/astropy_xarray/formatting.py b/astropy_xarray/formatting.py index 801c451d..b317b1e8 100644 --- a/astropy_xarray/formatting.py +++ b/astropy_xarray/formatting.py @@ -90,6 +90,8 @@ def last_n_items(array, n_desired): # based on xarray.core.formatting.format_item def format_item(x, quote_strings=True): """Returns a succinct summary of an object as a string""" + if isinstance(x, np.str_): + return repr(str(x)) if quote_strings else x if isinstance(x, (str, bytes)): return repr(x) if quote_strings else x elif isinstance(x, float): diff --git a/astropy_xarray/tests/test_accessors_coords.py b/astropy_xarray/tests/test_accessors_coords.py new file mode 100644 index 00000000..2375b0b2 --- /dev/null +++ b/astropy_xarray/tests/test_accessors_coords.py @@ -0,0 +1,117 @@ +import astropy.units as u +import numpy as np +import pytest +import xarray as xr +from astropy.coordinates import Distance, Latitude, Longitude + +import astropy_xarray # noqa: F401 + + +@pytest.mark.parametrize("unit", ["rad", "deg", "arcmin", "arcsec"]) +@pytest.mark.parametrize("class_type", [Latitude, Longitude]) +def test_angle_dataarray_dequantify(class_type, unit): + da: xr.DataArray = xr.DataArray( + class_type([0, 0.5, 1], unit=unit) + ).astropy.dequantify() + + assert da.attrs == {"units": {"class": class_type.__name__.lower(), "unit": unit}} + np.testing.assert_array_equal( + da.data, class_type([0, 0.5, 1], unit=unit).value, strict=True + ) + + +@pytest.mark.parametrize("unit", ["m", "km", "Mpc"]) +@pytest.mark.parametrize("class_type", [Distance]) +def test_distance_dataarray_dequantify(class_type, unit): + # TODO: support allow_negative=True + da: xr.DataArray = xr.DataArray( + class_type([0, 1, 2], unit=unit) + ).astropy.dequantify() + + assert da.attrs == {"units": {"class": class_type.__name__.lower(), "unit": unit}} + np.testing.assert_array_equal( + da.data, class_type([0, 1, 2], unit=unit).value, strict=True + ) + + +@pytest.mark.parametrize("unit", ["rad", "deg", "arcmin", "arcsec"]) +@pytest.mark.parametrize("class_type", [Latitude, Longitude]) +def test_coord_dataarray_quantify(class_type, unit): + da = xr.DataArray( + class_type([0, 0.5, 1], unit=unit).value, + attrs={"units": {"class": class_type.__name__.lower(), "unit": unit}}, + ).astropy.quantify() + + assert isinstance(da.data, class_type) + np.testing.assert_array_equal( + # da, + da.data, + class_type([0, 0.5, 1], unit=unit), + strict=True, + ) + + +def test_coord_dataset_dequantify(): + ds = xr.Dataset( + { + "lat": ("time", Latitude([0, 1, 2] * u.deg)), + "long": ("time", Longitude([0, 1, 2] * u.deg)), + "dist": ("time", Distance([0, 1, 2] * u.Mpc)), + "temp": ("time", u.Quantity([0, 1, 2], u.C)), + } + ).astropy.dequantify() + + assert ds.lat.attrs == {"units": {"class": "latitude", "unit": "deg"}} + np.testing.assert_array_equal( + ds.lat, Latitude([0, 1, 2], unit="deg").value, strict=True + ) + + assert ds.long.attrs == {"units": {"class": "longitude", "unit": "deg"}} + np.testing.assert_array_equal( + ds.lat, Longitude([0, 1, 2], unit="deg").value, strict=True + ) + + assert ds.dist.attrs == {"units": {"class": "distance", "unit": "Mpc"}} + np.testing.assert_array_equal( + ds.lat, Longitude([0, 1, 2], unit="deg").value, strict=True + ) + + +def test_coord_dataset_quantify(): + ds = xr.Dataset( + { + "lat": xr.DataArray( + Latitude([0, 1, 2] * u.deg).value, + dims="time", + attrs={"units": {"class": "latitude", "unit": "deg"}}, + ), + "long": xr.DataArray( + Longitude([0, 1, 2] * u.deg).value, + dims="time", + attrs={"units": {"class": "longitude", "unit": "deg"}}, + ), + "dist": xr.DataArray( + Distance([0, 1, 2], "Mpc").value, + dims="time", + attrs={"units": {"class": "distance", "unit": "Mpc"}}, + ), + "temp": xr.DataArray( + u.Quantity([0, 1, 2], "K").value, dims="time", attrs={"units": "K"} + ), + } + ).astropy.quantify() + + assert isinstance(ds.lat.data, Latitude) + np.testing.assert_array_equal( + ds.lat.data, Latitude([0, 1, 2], unit="deg"), strict=True + ) + + assert isinstance(ds.long.data, Longitude) + np.testing.assert_array_equal( + ds.long.data, Longitude([0, 1, 2], unit="deg"), strict=True + ) + + assert isinstance(ds.dist.data, Distance) + np.testing.assert_array_equal( + ds.dist.data, Distance([0, 1, 2], unit="Mpc"), strict=True + ) diff --git a/astropy_xarray/tests/test_accessors_log.py b/astropy_xarray/tests/test_accessors_log.py new file mode 100644 index 00000000..3d4d4991 --- /dev/null +++ b/astropy_xarray/tests/test_accessors_log.py @@ -0,0 +1,104 @@ +import astropy.units as u +import numpy as np +import pytest +import xarray as xr + +import astropy_xarray # noqa: F401 + +# referenceless +# y(x) dB = 10log_10(x) -> x = 10^(y/10) +# y(x) mag = -2.5log_10(x) -> x = 10^(-y/2.5) +# y(x) dex = log_10(x) -> x = 10^(y) + +# with reference +# y dB(2 V) = 10log_10(x/2 V) -> x = 2 V * 10^(y/10) +# y mag(2 V) = -2.5log_10(x/2 V) -> x = 2 V * 10^(-y/2.5) +# y dex(2 V) = log_10(x/2 V) -> x = 2 V * 10^(y) + + +@pytest.mark.parametrize("unit", ["dex", "dB", "mag", "dex(cd)", "dB(cd)", "mag(cd)"]) +def test_log_dataarray_dequantify(unit): + da: xr.DataArray = xr.DataArray( + u.LogQuantity([0, 0.5, 1], unit=unit) + ).astropy.dequantify() + + assert da.attrs == {"units": unit} + np.testing.assert_array_equal( + da.data, u.LogQuantity([0, 0.5, 1], unit=unit).value, strict=True + ) + + +@pytest.mark.parametrize("unit", ["dex(10 cd)", "dB(10 cd)", "mag(10 cd)"]) +def test_log_dataarray_quantify(unit): + da: xr.DataArray = xr.DataArray( + u.LogQuantity([0, 0.5, 1], unit=unit).value, attrs={"units": unit} + ).astropy.quantify() + + assert isinstance(da.data, u.LogQuantity) + np.testing.assert_array_equal( + da.data, u.LogQuantity([0, 0.5, 1], unit=unit), strict=True + ) + + +def test_log_dataset_dequantify(): + ds: xr.DataArray = xr.Dataset( + { + "volume": ("time", [0, 1, 2] * u.dB("20 uPa")), + "brightness": ("time", [0, 1, 2] * u.mag("cd")), + "metallicity": ("time", [0, 1, 2] * u.dex), + "temp": ("time", u.Quantity([0, 1, 2], u.C)), + } + ).astropy.dequantify() + + assert ds.volume.attrs == {"units": "dB(20 uPa)"} + np.testing.assert_array_equal(ds.volume, np.array([0.0, 1.0, 2.0]), strict=True) + + assert ds.brightness.attrs == {"units": "mag(cd)"} + np.testing.assert_array_equal(ds.brightness, np.array([0.0, 1.0, 2.0]), strict=True) + + assert ds.metallicity.attrs == {"units": "dex"} + np.testing.assert_array_equal( + ds.metallicity, np.array([0.0, 1.0, 2.0]), strict=True + ) + + +def test_coord_dataset_quantify(): + ds: xr.DataArray = xr.Dataset( + { + "volume": xr.DataArray( + np.array([0.0, 1.0, 2.0]), dims="time", attrs={"units": "dB(20 uPa)"} + ), + "brightness": xr.DataArray( + np.array([0.0, 1.0, 2.0]), dims="time", attrs={"units": "mag(cd)"} + ), + "metallicity": xr.DataArray( + np.array([0.0, 1.0, 2.0]), dims="time", attrs={"units": "dex"} + ), + "temp": xr.DataArray( + np.array([0.0, 1.0, 2.0]), dims="time", attrs={"units": "K"} + ), + } + ).astropy.quantify() + + assert isinstance(ds.volume.data, u.LogQuantity) + np.testing.assert_array_equal( + ds.volume.data, [0, 1, 2] * u.dB("20 uPa"), strict=True + ) + np.testing.assert_array_equal( + ds.volume.data.to("uPa"), + [20.0, 25.178508235883346, 31.697863849222273] * u.uPa, + strict=True, + ) + + assert isinstance(ds.brightness.data, u.LogQuantity) + np.testing.assert_array_equal( + ds.brightness.data, [0, 1, 2] * u.mag("cd"), strict=True + ) + np.testing.assert_array_equal( + ds.brightness.data.to("cd"), + [1.0, 0.3981071705534972, 0.15848931924611134] * u.cd, + strict=True, + ) + + assert isinstance(ds.metallicity.data, u.Quantity) + np.testing.assert_array_equal(ds.metallicity.data, [0, 1, 2] * u.dex, strict=True) diff --git a/astropy_xarray/tests/test_accessors_time.py b/astropy_xarray/tests/test_accessors_time.py new file mode 100644 index 00000000..6dd72114 --- /dev/null +++ b/astropy_xarray/tests/test_accessors_time.py @@ -0,0 +1,118 @@ +import astropy.units as u +import numpy as np +import pytest +import xarray as xr +from astropy.time import Time, TimeDelta + +import astropy_xarray # noqa: F401 + + +@pytest.mark.parametrize("class_type", [Time, TimeDelta]) +def test_time_dataarray_dequantify(class_type): + da: xr.DataArray = xr.DataArray( + class_type([0, 1, 2], format="jd", scale="ut1") + ).astropy.dequantify() + + assert da.attrs == { + "units": { + "class": class_type.__name__.lower(), + "format": "jd", + "precision": 3, + "scale": "ut1", + } + } + np.testing.assert_array_equal( + da.data, class_type([0, 1, 2], format="jd", scale="ut1").value, strict=True + ) + + +@pytest.mark.parametrize("class_type", [Time, TimeDelta]) +def test_time_dataarray_quantify(class_type): + da = xr.DataArray( + class_type([0, 1, 2], format="jd", scale="ut1").value, + attrs={ + "units": { + "class": class_type.__name__.lower(), + "format": "jd", + "precision": 3, + "scale": "ut1", + } + }, + ).astropy.quantify() + + assert isinstance(da.data, class_type) + np.testing.assert_array_equal( + da, class_type([0, 1, 2], format="jd", scale="ut1"), strict=True + ) + + +def test_time_dataset_dequantify(): + ds = xr.Dataset( + { + "timestamps": ("time", Time([0, 1, 2], format="jd", scale="ut1")), + "timesteps": ("time", TimeDelta([0, 1, 2], format="jd", scale="ut1")), + "secs": ("time", u.Quantity([0, 1, 2], "s")), + "temp": ("time", u.Quantity([0, 1, 2], "K")), + } + ).astropy.dequantify() + + assert ds.timestamps.attrs == { + "units": {"class": "time", "format": "jd", "precision": 3, "scale": "ut1"} + } + np.testing.assert_array_equal( + ds.timestamps, Time([0, 1, 2], format="jd", scale="ut1").value, strict=True + ) + + assert ds.timesteps.attrs == { + "units": {"class": "timedelta", "format": "jd", "precision": 3, "scale": "ut1"} + } + np.testing.assert_array_equal( + ds.timesteps, TimeDelta([0, 1, 2], format="jd", scale="ut1").value, strict=True + ) + + +def test_time_dataset_quantify(): + ds = xr.Dataset( + { + "timestamps": xr.DataArray( + Time([0, 1, 2], format="jd", scale="ut1").value, + dims="time", + attrs={ + "units": { + "class": "time", + "format": "jd", + "precision": 3, + "scale": "ut1", + } + }, + ), + "timesteps": xr.DataArray( + TimeDelta([0, 1, 2], format="jd", scale="ut1").value, + dims="time", + attrs={ + "units": { + "class": "timedelta", + "format": "jd", + "precision": 3, + "scale": "ut1", + } + }, + ), + "secs": xr.DataArray( + u.Quantity([0, 1, 2], "s").value, dims="time", attrs={"units": "s"} + ), + "temp": xr.DataArray( + u.Quantity([0, 1, 2], "K").value, dims="time", attrs={"units": "K"} + ), + } + ).astropy.quantify() + + assert isinstance(ds.timestamps.data, Time) + np.testing.assert_array_equal( + ds.timestamps, Time([0, 1, 2], format="jd", scale="ut1"), strict=True + ) + + assert isinstance(ds.timesteps.data, TimeDelta) + np.testing.assert_array_equal( + ds.timesteps, TimeDelta([0, 1, 2], format="jd", scale="ut1"), strict=True + ) diff --git a/astropy_xarray/tests/test_frame_conversion.py b/astropy_xarray/tests/test_frame_conversion.py new file mode 100644 index 00000000..77945eaf --- /dev/null +++ b/astropy_xarray/tests/test_frame_conversion.py @@ -0,0 +1,85 @@ +import astropy.units as u +import numpy.testing +from astropy.coordinates import ( + CartesianDifferential, + CartesianRepresentation, + SkyCoord, +) +from astropy.coordinates.builtin_frames import ICRS + +from astropy_xarray.coordinates import ( + skycoord_to_dataset, +) +from astropy_xarray.coordinates.sky_coord import ( + _skycoord_component_names, +) + + +def test_frame_conversion_aliasing(): + sc = SkyCoord( + ICRS().realize_frame( + CartesianRepresentation( + [0.1, 0.2] * u.km, + [0.5, 0.7] * u.km, + [0.3, 0.4] * u.km, + differentials=CartesianDifferential( + d_x=[0.1, 0.2] * u.km / u.s, + d_y=[0.5, 0.7] * u.km / u.s, + d_z=[0.3, 0.4] * u.km / u.s, + ), + ) + ) + ) + sc2 = SkyCoord( + ICRS().realize_frame( + sc.represent_as( + sc.frame.default_representation, sc.frame.default_differential + ) + ) + ) + numpy.testing.assert_array_equal(sc.pm_ra_cosdec, sc2.pm_ra_cosdec) + + +def test_frame_conversion_components(): + sc = SkyCoord( + ICRS().realize_frame( + CartesianRepresentation( + [[0.1], [0.2]] * u.km, + [[0.5], [0.7]] * u.km, + [[0.3], [0.4]] * u.km, + differentials=CartesianDifferential( + d_x=[[0.1], [0.2]] * u.km / u.s, + d_y=[[0.5], [0.7]] * u.km / u.s, + d_z=[[0.3], [0.4]] * u.km / u.s, + ), + ) + ) + ) + + assert list(_skycoord_component_names(sc)) == [ + "x", + "y", + "z", + "v_x", + "v_y", + "v_z", + ] + + skycoord_to_dataset(sc, None) + + sc2 = SkyCoord( + ICRS().realize_frame( + sc.represent_as( + sc.frame.default_representation, sc.frame.default_differential + ) + ) + ) + + assert list(_skycoord_component_names(sc2)) == [ + "ra", + "dec", + "distance", + "pm_ra_cosdec", + "pm_dec", + "radial_velocity", + ] diff --git a/astropy_xarray/tests/test_sky_coord.py b/astropy_xarray/tests/test_sky_coord.py new file mode 100644 index 00000000..b0c9d42f --- /dev/null +++ b/astropy_xarray/tests/test_sky_coord.py @@ -0,0 +1,348 @@ +import astropy.units as u +import numpy as np +import pytest +import xarray as xr +from astropy.coordinates import ( + Angle, + BaseCoordinateFrame, + BaseRepresentation, + CartesianDifferential, + CartesianRepresentation, + CylindricalDifferential, + CylindricalRepresentation, + Distance, + EarthLocation, + Latitude, + Longitude, + SkyCoord, + SphericalCosLatDifferential, + SphericalRepresentation, + UnitSphericalCosLatDifferential, + UnitSphericalDifferential, + UnitSphericalRepresentation, +) +from astropy.coordinates.builtin_frames import ( + CIRS, + FK4, + FK5, + GCRS, + HCRS, + ICRS, + ITRS, + AltAz, + CustomBarycentricEcliptic, + Galactic, + GalacticLSR, + Galactocentric, + GeocentricMeanEcliptic, + GeocentricTrueEcliptic, + HADec, + HeliocentricEclipticIAU76, + HeliocentricMeanEcliptic, + HeliocentricTrueEcliptic, + PrecessedGeocentric, + Supergalactic, +) +from astropy.time import Time + +from astropy_xarray.coordinates import ( + dataset_to_skycoord, + skycoord_to_dataset, +) +from astropy_xarray.coordinates.sky_coord import ( + _skycoord_differential_component_names, + _skycoord_representation_component_names, +) + + +@pytest.mark.parametrize( + ("expected_base", "expected_s"), + [ + ( + ("ra", "dec"), + ("pm_ra_cosdec", "pm_dec", "radial_velocity"), + ), + ], +) +def test_unitspherical_repr_components(expected_base, expected_s): + sc = SkyCoord( + ra=[[0.1], [2], [0.2]] * u.deg, + dec=[[0.5], [7], [0.7]] * u.deg, + pm_ra_cosdec=[[0.002], [0.002], [0.002]] * u.deg / u.yr, + pm_dec=[[0.002], [0.002], [0.002]] * u.deg / u.yr, + radial_velocity=[[0.0], [0.0], [0.0]] * u.deg / u.yr, + ) + expected_repr = str(sc) + assert _skycoord_representation_component_names(sc) == expected_base + assert _skycoord_differential_component_names(sc) == expected_s + assert str(sc) == expected_repr + + +@pytest.mark.parametrize( + ("expected_base", "expected_s"), + [ + ( + ("ra", "dec", "distance"), + ("pm_ra_cosdec", "pm_dec"), + ) + ], +) +def test_unitspherical_diff_components(expected_base, expected_s): + sc = SkyCoord( + ra=[[0.1], [2], [0.2]] * u.deg, + dec=[[0.5], [7], [0.7]] * u.deg, + distance=[[1.0], [1.0], [1.0]] * u.dimensionless_unscaled, + pm_ra_cosdec=[[0.002], [0.002], [0.002]] * u.deg / u.yr, + pm_dec=[[0.002], [0.002], [0.002]] * u.deg / u.yr, + ) + expected_str = str(sc) + assert _skycoord_representation_component_names(sc) == expected_base + assert _skycoord_differential_component_names(sc) == expected_s + assert str(sc) == expected_str + + +@pytest.mark.parametrize( + ("frame", "expected_base", "expected_s"), + [ + ( + ICRS(), + ("ra", "dec"), + ("pm_ra_cosdec", "pm_dec"), + ), + ( + AltAz(), + ("az", "alt"), + ("pm_az_cosalt", "pm_alt"), + ), + ( + ITRS(), + ("lon", "lat"), + ("pm_lon_coslat", "pm_lat"), + ), + ], +) +def test_spherical_direction_components_realize_frame(frame, expected_base, expected_s): + rep = UnitSphericalRepresentation( + [[0.1], [2], [0.2]] * u.deg, + [[0.5], [7], [0.7]] * u.deg, + differentials=UnitSphericalCosLatDifferential( + [[0.002], [0.002], [0.002]] * u.deg / u.yr, + [[0.002], [0.002], [0.002]] * u.deg / u.yr, + ), + ) + sc = SkyCoord(frame.realize_frame(rep)) + assert _skycoord_representation_component_names(sc) == expected_base + assert _skycoord_differential_component_names(sc) == expected_s + + +@pytest.mark.parametrize( + ("expected_base", "expected_s"), + [ + (("x", "y", "z"), ("v_x", "v_y", "v_z")), + ], +) +def test_cartesian_direction_components(expected_base, expected_s): + sc = SkyCoord( + x=[[0.1], [2], [0.2]] * u.dimensionless_unscaled, + y=[[0.5], [7], [0.7]] * u.dimensionless_unscaled, + z=[[0.5], [7], [0.7]] * u.dimensionless_unscaled, + v_x=[[0.002], [0.002], [0.002]] * u.dimensionless_unscaled / u.yr, + v_y=[[0.002], [0.002], [0.002]] * u.dimensionless_unscaled / u.yr, + v_z=[[0.002], [0.002], [0.002]] * u.dimensionless_unscaled / u.yr, + frame="itrs", + ) + assert _skycoord_representation_component_names(sc) == expected_base + assert _skycoord_differential_component_names(sc) == expected_s + + +@pytest.mark.parametrize( + "use_default_representation", + [ + pytest.param(True), + pytest.param(False), + ], +) +@pytest.mark.parametrize( + ("expected_data_classes", "data"), + [ + pytest.param( + (Longitude, Latitude), + UnitSphericalRepresentation( + [[0.1], [0.2]] * u.deg, + [[0.5], [0.7]] * u.deg, + differentials=UnitSphericalDifferential( + d_lon=[[0.001], [0.002]] * u.mas / u.yr, + d_lat=[[0.001], [0.002]] * u.mas / u.yr, + ), + ), + id="sphericalcoslat-2", + ), + pytest.param( + (Longitude, Latitude), + UnitSphericalRepresentation( + [[0.1], [0.2]] * u.deg, + [[0.5], [0.7]] * u.deg, + differentials=UnitSphericalCosLatDifferential( + d_lon_coslat=[[0.001], [0.002]] * u.mas / u.yr, + d_lat=[[0.001], [0.002]] * u.mas / u.yr, + ), + ), + id="spherical-2", + ), + pytest.param( + (Longitude, Latitude, Distance), + SphericalRepresentation( + [[0.1], [0.2]] * u.deg, + [[0.5], [0.7]] * u.deg, + [[0.3], [0.4]] * u.pc, + differentials=SphericalCosLatDifferential( + d_lon_coslat=[[0.001], [0.002]] * u.mas / u.yr, + d_lat=[[0.001], [0.002]] * u.mas / u.yr, + d_distance=[[0.0002], [0.0002]] * u.km / u.s, + ), + ), + id="spherical-3", + ), + pytest.param( + (u.Quantity, u.Quantity, u.Quantity), + CartesianRepresentation( + [[0.1], [0.2]] * u.km, + [[0.5], [0.7]] * u.km, + [[0.3], [0.4]] * u.km, + differentials=CartesianDifferential( + d_x=[[0.1], [0.2]] * u.km / u.s, + d_y=[[0.5], [0.7]] * u.km / u.s, + d_z=[[0.3], [0.4]] * u.km / u.s, + ), + ), + id="cartesian-3", + ), + pytest.param( + (u.Quantity, Angle, u.Quantity), + CylindricalRepresentation( + [[0.1], [0.2]] * u.km, + [[0.5], [0.7]] * u.deg, + [[0.3], [0.4]] * u.km, + differentials=CylindricalDifferential( + d_rho=[[0.1], [0.2]] * u.km / u.s, + d_phi=[[0.5], [0.7]] * u.deg / u.s, + d_z=[[0.3], [0.4]] * u.km / u.s, + ), + ), + id="cylindrical-3", + ), + ], +) +@pytest.mark.parametrize( + "frame", + [ + ICRS(), + FK5(), + FK4(), + GCRS(), + GCRS(obstime=Time("2024-01-01")), + CIRS(), + HCRS(), + ITRS(), + AltAz(), + AltAz(obstime=Time("2025-06-01"), location=EarthLocation.of_site("greenwich")), + HADec(), + HADec(obstime=Time("2025-06-01"), location=EarthLocation.of_site("greenwich")), + Galactic(), + GalacticLSR(), + PrecessedGeocentric(), + GeocentricMeanEcliptic(), + GeocentricTrueEcliptic(), + HeliocentricMeanEcliptic(), + HeliocentricTrueEcliptic(), + HeliocentricEclipticIAU76(), + Galactocentric(), + Galactocentric( + galcen_distance=3 * u.kpc, + galcen_coord=ICRS( + ra=1 * u.deg, + dec=1 * u.deg, + pm_ra_cosdec=0.1 * u.mas / u.yr, + pm_dec=0.1 * u.mas / u.yr, + ), + galcen_v_sun=CartesianDifferential( + d_x=1 * u.km / u.s, d_y=1 * u.km / u.s, d_z=1 * u.km / u.s + ), + z_sun=u.Quantity(30 * u.pc), + roll=u.Quantity(20 * u.deg), + ), + Supergalactic(), + CustomBarycentricEcliptic(), + CustomBarycentricEcliptic(obliquity=1 * u.deg), + ], + ids=lambda param: param.name, +) +def test_skycoord_roundtrip( + frame: BaseCoordinateFrame, + data: BaseRepresentation, + expected_data_classes: tuple[type, ...], + use_default_representation: bool, +): + # all frames support all representation types + representation_type = type(data) + differential_type = type(data.differentials["s"]) + + if use_default_representation: + expected = SkyCoord( + frame.realize_frame( + data.represent_as( + frame.default_representation, frame.default_differential + ) + ) + ) + expected_frame_data_rep_name = frame.default_representation.get_name() + expected_frame_data_diff_name = frame.default_differential.get_name() + else: + expected = SkyCoord(frame.realize_frame(data)) + expected_frame_data_rep_name = representation_type.get_name() + expected_frame_data_diff_name = differential_type.get_name() + + assert expected.frame.data.get_name() == expected_frame_data_rep_name + assert ( + expected.frame.data.differentials["s"].get_name() + == expected_frame_data_diff_name + ) + + coords = [ + ("calibrator_name", ["a1", "a2"]), + ("time_polynomial", [0]), + ] + ds = skycoord_to_dataset(expected, coords=coords) + + # data_var keys + expected_keys_frame = expected.frame.copy() + expected_keys_frame.representation_type = type(expected.frame.data) + expected_keys_frame.differential_type = type( + expected.frame.data.differentials.get("s") + ) + expected_keys = set(expected_keys_frame.get_representation_component_names()) | set( + expected_keys_frame.get_representation_component_names("s") + ) + assert set(ds.data_vars) == expected_keys + + # Convert back + actual = dataset_to_skycoord(ds) + + # Check Representation + for component_name in actual.frame.representation_component_names: + np.testing.assert_array_equal( + getattr(actual, component_name), getattr(expected, component_name) + ) + + # Check Differentials + assert expected.data.differentials["s"] is not None + np.testing.assert_array_equal( + actual.data.differentials["s"], expected.data.differentials["s"], strict=True + ) + + # Check SkyCoords equal + assert actual.representation_type == expected.representation_type + assert actual.differential_type == expected.differential_type + assert actual.is_equivalent_frame(expected) + assert ds.coords.equals(xr.Coordinates(dict(coords))) + np.testing.assert_array_equal(actual, expected) diff --git a/astropy_xarray/tests/test_time.py b/astropy_xarray/tests/test_time.py new file mode 100644 index 00000000..ee690753 --- /dev/null +++ b/astropy_xarray/tests/test_time.py @@ -0,0 +1,117 @@ +import numpy as np +import pytest +from astropy.time import Time + +import astropy_xarray # noqa: F401 + + +@pytest.fixture(name="time") +def time_fixture(value, format, scale): + return Time(value, format=format, scale=scale) + + +@pytest.mark.parametrize("scale", ["utc", "tai", "ut1"]) +@pytest.mark.parametrize( + "value,format,dtype", + [ + (np.array([1.0, 2.0, 3.0]), "unix", "float64"), + (np.array([1.0, 2.0, 3.0]), "unix_tai", "float64"), + (np.array([1.0, 2.0, 3.0]), "mjd", "float64"), + ( + np.array( + [ + np.datetime64("1970-01-01T00:00:09.000082030"), + np.datetime64("1970-01-01T00:00:10.000082060"), + np.datetime64("1970-01-01T00:00:11.000082090"), + ] + ), + "datetime64", + " bytes: + if isinstance(data, xr.DataTree): + dict_data = { + group: ds.to_dict("array") + for group, ds in data.astropy.dequantify().to_dict().items() + } + else: + dict_data = data.astropy.dequantify().to_dict("array") + return json.dumps(dict_data) # type: ignore + + +def astropy_encode_msgpack(data: xr.DataArray | xr.Dataset | xr.DataTree) -> bytes: + if isinstance(data, xr.DataTree): + dict_data = { + group: ds.to_dict("array") + for group, ds in data.astropy.dequantify().to_dict().items() + } + else: + dict_data = data.astropy.dequantify().to_dict("array") + return msgpack.packb(dict_data, default=msgpack_numpy.encode) # type: ignore + + +def astropy_decode_msgpack(data: bytes | memoryview): + data_dict: dict = msgpack.unpackb(data, object_hook=msgpack_numpy.decode) + if "/" in data_dict: + dt_dict = { + group: xr.Dataset.from_dict(ds_dict).astropy.quantify() + for group, ds_dict in data_dict.items() + } + return xr.DataTree.from_dict(dt_dict) + if "data_vars" in data_dict: + return xr.Dataset.from_dict(data_dict).astropy.quantify() + else: + return xr.DataArray.from_dict(data_dict).astropy.quantify() + + +def generate_baselines(antenna_count) -> list[tuple[int, int]]: + baselines = [] + for i in range(antenna_count): + for j in range(i + 1, antenna_count): + baselines.append((i, j)) + return baselines + + +def test_simple_visibility_dataset(): + a = 4 # antennas + b = int(a * (a - 1) / 2) # baselines + p = 4 # polarisations + f = 7 # frequencies + t = 2 # timestamps + + # import cf_xarray as cfxr + baselines = generate_baselines(a) + antenna1, antenna2 = map(np.array, zip(*baselines)) + + VISIBILITY = xr.DataTree( + xr.Dataset( + coords=dict( + time=xr.DataArray( + [1.0, 2.0] * u.s, + dims=["time"], + ), + frequency=xr.DataArray( + np.arange(1000.0, 1007.0) * u.Hz, + dims=["frequency"], + ), + polarisation=xr.DataArray( + ["XX", "XY", "YX", "YY"], dims=["polarisation"] + ), + # simple index + baseline=xr.DataArray(np.arange(len(baselines)), dims=["baseline"]), + antenna1=xr.DataArray(antenna1, dims=["baseline"]), + antenna2=xr.DataArray(antenna2, dims=["baseline"]), + uvw_label=xr.DataArray(["u", "v", "w"], dims=["uvw_label"]), + ), + data_vars=dict( + uvw=xr.DataArray( + data=[ + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 0.0], + [0.5, 0.5, 0.5], + [0.0, 0.0, 0.0], + ], + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 0.0], + [0.5, 0.5, 0.5], + [0.0, 0.0, 0.0], + ], + ] + * u.m, + dims=["time", "baseline", "uvw_label"], + ), + vis=xr.DataArray( + data=np.ones([t, b, f, p]), + dims=["time", "baseline", "frequency", "polarisation"], + ), + ), + ), + children=dict( + field_phase_center=xr.DataTree( + skycoord_to_dataset( + SkyCoord(ra=[0.1] * u.deg, dec=[0.5] * u.deg, frame=ICRS()), + coords=[ + ("time_poly", [0]), + ], + ) + ), + # extensions + calibrator=xr.DataTree( + skycoord_to_dataset( + SkyCoord( + ra=[[0.5], [1.5], [0.2], [0.9], [1.1]] * u.deg, + dec=[[1.2], [0.8], [0.6], [1.0], [1.5]] * u.deg, + ), + coords=[ + ("calibrator_label", ["A", "B", "C", "D", "E"]), + ("time_poly", [0]), + ], + ) + ), + ), + ) + + message = astropy_encode_msgpack(VISIBILITY) + assert isinstance(message, bytes) + ds_out = astropy_decode_msgpack(message) + xr.testing.assert_identical(VISIBILITY, ds_out) + + assert dataset_to_skycoord(ds_out["field_phase_center"]) == ( + SkyCoord(ra=[0.1] * u.deg, dec=[0.5] * u.deg, frame=ICRS()) + ) + + +@pytest.mark.parametrize( + ("skycoords", "frame", "unit"), + [ + ( + FIELD_PHASE_CENTER_ICRS, + {"name": "icrs"}, + {"ra": "deg", "dec": "deg"}, + ), + ( + FIELD_PHASE_CENTER_FK5, + {"name": "fk5", "equinox": 1}, + {"ra": "deg", "dec": "deg"}, + ), + (FIELD_PHASE_CENTER_GEO, {"name": "gcrs"}, {"x": "km", "y": "km", "z": "km"}), + ], +) +def test_visibility_dataset(skycoords: list[SkyCoord], frame, unit): + a = 4 # antennas + b = int(a * (a - 1) / 2) # baselines + p = 4 # polarisations + f = 7 # frequencies + t = 2 # timestamps + + VISIBILITY = xr.DataTree( + xr.Dataset( + coords=dict( + time=xr.DataArray( + [1.0, 2.0] * u.s, + dims=["time"], + ), + frequency=xr.DataArray( + np.arange(1000.0, 1007.0) * u.Hz, + dims=["frequency"], + ), + polarisation=xr.DataArray( + ["XX", "XY", "YX", "YY"], dims=["polarisation"] + ), + uvw_label=xr.DataArray(["u", "v", "w"], dims=["uvw_label"]), + ) + | dict( + xr.Coordinates.from_pandas_multiindex( + pd.MultiIndex.from_tuples( + generate_baselines(a), names=("antenna1", "antenna2") + ), + dim="baseline", + ) + ), + data_vars=dict( + uvw=xr.DataArray( + data=[ + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 0.0], + [0.5, 0.5, 0.5], + [0.0, 0.0, 0.0], + ], + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 0.0], + [0.5, 0.5, 0.5], + [0.0, 0.0, 0.0], + ], + ] + * u.m, + dims=["time", "baseline", "uvw_label"], + ), + vis=xr.DataArray( + data=np.ones([t, b, f, p]), + dims=["time", "baseline", "frequency", "polarisation"], + ), + ), + ), + children=dict( + field_phase_center=xr.DataTree( + skycoord_to_dataset( + SkyCoord( + ra=[[0.1, 0.2]] * u.deg, + dec=[[0.5, 0.2]] * u.deg, + distance=[[0.0, 0.0]] * u.km, + pm_ra_cosdec=[[1.0, 1.0]] * u.deg / u.s, + pm_dec=[[0.5, 0.2]] * u.deg / u.s, + radial_velocity=[[1, 1]] * u.pc / u.yr, + frame=ICRS(), + ), + coords=[ + ("phase_center_label", np.array([0])), + ("time_poly", np.array([0, 1])), + ], + ) + ), + # extensions + calibrator=xr.DataTree( + skycoord_to_dataset( + SkyCoord( + ra=[[0.5], [1.5], [0.2], [0.9], [1.1]] * u.deg, + dec=[[1.2], [0.8], [0.6], [1.0], [1.5]] * u.deg, + ), + coords=[ + ("calibrator_label", ["A", "B", "C", "D", "E"]), + ("time_poly", np.array([0])), + ], + ) + ), + ), + ) + + m = astropy_encode_msgpack(VISIBILITY) + ds_out = astropy_decode_msgpack(m) + xr.testing.assert_identical(VISIBILITY, ds_out) diff --git a/astropy_xarray/time_compat.py b/astropy_xarray/time_compat.py new file mode 100644 index 00000000..ade4f0ce --- /dev/null +++ b/astropy_xarray/time_compat.py @@ -0,0 +1,55 @@ +import numpy as np +from astropy.time import Time, TimeDelta + +import astropy_xarray +import astropy_xarray.formatting + + +def time_inline_repr(time: Time, max_width: int): + value = time.datetime64 + scale_repr = f"{time.format} {time.scale}" + if isinstance(value, np.ndarray): + data_repr = astropy_xarray.formatting.format_array_flat( + value, max_width - len(scale_repr) - 3 + ) + else: + data_repr = astropy_xarray.formatting.maybe_truncate( + repr(value), max_width - len(scale_repr) - 3 + ) + + return f"[{scale_repr}] {data_repr}" + + +def time_delta_inline_repr(time: TimeDelta, max_width: int): + value = time.value + scale_repr = f"{time.format} {time.scale}" + if isinstance(value, np.ndarray): + data_repr = astropy_xarray.formatting.format_array_flat( + value, max_width - len(scale_repr) - 3 + ) + else: + data_repr = astropy_xarray.formatting.maybe_truncate( + repr(value), max_width - len(scale_repr) - 3 + ) + + return f"[{scale_repr}] {data_repr}" + + +Time.dtype = property(lambda self: self.value.dtype) +Time.nbytes = property(lambda self: self.value.nbytes) +Time.__array__ = lambda obj: obj.value +Time.__array_ufunc__ = lambda obj, *args, **kwargs: obj.value.__array_ufunc__( + *args, **kwargs +) +Time._repr_inline_ = time_inline_repr +Time.__array_namespace = np + + +TimeDelta.dtype = property(lambda self: self.value.dtype) +TimeDelta.nbytes = property(lambda self: self.value.nbytes) +TimeDelta.__array__ = lambda obj: obj.value +TimeDelta.__array_ufunc__ = lambda obj, *args, **kwargs: obj.value.__array_ufunc__( + *args, **kwargs +) +TimeDelta._repr_inline_ = time_delta_inline_repr +TimeDelta.__array_namespace = np diff --git a/docs/api.rst b/docs/api.rst index 0c3ec879..4b36f83d 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -5,6 +5,11 @@ This page contains a auto-generated summary of ``astropy-xarray``'s API. .. autosummary:: :toctree: generated/ + astropy_xarray.coordinates.dataset_to_skycoord + astropy_xarray.coordinates.skycoord_to_dataset + astropy_xarray.coordinates.load_frame + astropy_xarray.coordinates.load_representation + Dataset ------- .. autosummary:: diff --git a/docs/examples/plotting.ipynb b/docs/examples/plotting.ipynb index 84578c3d..25d95db6 100644 --- a/docs/examples/plotting.ipynb +++ b/docs/examples/plotting.ipynb @@ -21,9 +21,11 @@ "import astropy_xarray # noqa: F401\n", "\n", "# to be able to read unit attributes following the CF conventions\n", + "# import cf_xarray.units # must be imported before pint_xarray\n", "u.set_enabled_aliases(\n", " {\n", " \"degK\": u.Kelvin,\n", + " \"K\": u.Kelvin,\n", " \"degrees_north\": u.deg,\n", " \"degrees_east\": u.deg,\n", " },\n", @@ -66,7 +68,7 @@ "metadata": {}, "source": [ "
\n", - "Note: this example uses the data provided by the xarray.tutorial functions. As such, the units attributes follow the CF conventions, which pint does not understand by default. To still be able to read them we are using the registry provided by cf-xarray.\n", + "Note: this example uses the data provided by the xarray.tutorial functions. As such, the units attributes follow the CF conventions, which astropy does not understand by default. To still be able to read them we are using the registry provided by cf-xarray.\n", "
" ] }, @@ -96,12 +98,7 @@ "metadata": {}, "outputs": [], "source": [ - "monthly_means = (\n", - " quantified.astropy.to(\"deg_C\", equivalencies=u.temperature())\n", - " .sel(time=\"2013\")\n", - " .groupby(\"time.month\")\n", - " .mean()\n", - ")\n", + "monthly_means = quantified.astropy.sel(time=\"2013\").groupby(\"time.month\").mean()\n", "monthly_means" ] }, @@ -146,10 +143,124 @@ "monthly_means.astropy.dequantify(format=\"unicode\").plot.imshow(col=\"month\", col_wrap=4)" ] }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "## Time\n", + "\n", + "astropy-xarray patches Time and DeltaTime to be recognized xarray, capturing:\n", + "* format\n", + "* timescale\n", + "* precision" + ] + }, { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "from astropy.time import Time, TimeDelta\n", + "\n", + "ds = xr.Dataset(\n", + " data_vars={\n", + " \"f_time\": (\"time_idx\", Time([0, 1, 2], format=\"unix\", scale=\"utc\")),\n", + " \"i_time\": (\n", + " \"time_idx\",\n", + " Time([0, 1, 2], format=\"unix\", scale=\"utc\").copy(\"datetime64\"),\n", + " ),\n", + " \"s_time\": (\n", + " \"time_idx\",\n", + " Time([0, 1, 2], format=\"unix\", scale=\"utc\").copy(\"isot\"),\n", + " ),\n", + " \"interval\": (\"time_idx\", TimeDelta([0.9, 0.9, 0.9], format=\"sec\", scale=\"tai\")),\n", + " \"temp\": (\"time_idx\", u.Quantity([0, 1, 2], \"K\")),\n", + " },\n", + " coords=xr.Coordinates(\n", + " {\n", + " \"time_idx\": [0, 1, 2],\n", + " }\n", + " ),\n", + ")\n", + "ds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "ds = ds.astropy.dequantify()\n", + "ds.i_time" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "## SkyCoord, Frame and Angle\n", + "\n", + "SkyCoord and Frame fully support conversion to datasets preserving:\n", + "* Representation\n", + "* Differential\n", + "* Frame\n", + "* Angle, Longitude, Latitude containers\n", + "\n", + "Conversion optionally requires coords argument for applying named coordinates to data variable axes.\n", + "\n", + "Frame-specific mapped names for data variables coming soon." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "from astropy.coordinates import SkyCoord\n", + "\n", + "from astropy_xarray.coordinates.sky_coord import (\n", + " dataset_to_skycoord,\n", + " skycoord_to_dataset,\n", + ")\n", + "\n", + "sc = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + ")\n", + "\n", + "\n", + "s = skycoord_to_dataset(sc, coords=[(\"field\", [0, 1, 2, 3])])\n", + "\n", + "# display(s1)\n", + "# display(s2)\n", + "s.astropy.dequantify()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "display(dataset_to_skycoord(s))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", "metadata": {}, "outputs": [], "source": [] diff --git a/docs/examples/requirements.txt b/docs/examples/requirements.txt new file mode 100644 index 00000000..8aa824de --- /dev/null +++ b/docs/examples/requirements.txt @@ -0,0 +1,3 @@ +scipy +pooch +matplotlib diff --git a/docs/examples/skycoord.ipynb b/docs/examples/skycoord.ipynb new file mode 100644 index 00000000..120e37a4 --- /dev/null +++ b/docs/examples/skycoord.ipynb @@ -0,0 +1,498 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## SkyCoord, Frame and Angles as xarray Dataset\n", + "\n", + "SkyCoord and Frame fully support conversion to datasets preserving:\n", + "* Representation\n", + "* Differential\n", + "* Frame\n", + "* Angle (Longitude and Latitude)\n", + "\n", + "Conversion to dataset supports optionally providing axis (xarray coordinate) labels." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import astropy.units as u\n", + "import numpy as np\n", + "import xarray as xr\n", + "from astropy.coordinates import ICRS, SkyCoord\n", + "\n", + "from astropy_xarray.coordinates.sky_coord import (\n", + " skycoord_to_dataset,\n", + ")\n", + "\n", + "sky_direction = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + " # NOTE: to explicity be a direction, use:\n", + " # frame=ICRS(\n", + " # representation_type='unitspherical',\n", + " # differential_type='unitsphericalcoslat',\n", + " # ),\n", + ")\n", + "display(sky_direction)\n", + "skycoord_to_dataset(sky_direction, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SkyCoordinate representation and differential types for both display and data are detected. In other words, radial_velocity when **implicit** is not recorded as a data_var.\n", + "\n", + "* *display* differential_type is `'sphericalcoslat'`, but\n", + "* *data* differential_type is `'unitsphericalcoslat'`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sky_position = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " distance=[4, 3, 6, 4] * u.parsec,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + ")\n", + "display(sky_position)\n", + "display(skycoord_to_dataset(sky_position, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])]))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Since the default ICRS differential type is 'sphericalcoslat', the first access\n", + "# to 'radial_velocity' will compute, cache and mutate the skycoordinate data\n", + "display(sky_position.radial_velocity)\n", + "\n", + "# skycoord_to_dataset will then subsequently include these radial velocity values when this occurs\n", + "display(skycoord_to_dataset(sky_position, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])]))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Convert Dataset to SkyCoord\n", + "\n", + "SkyCoord datasets use a frame attribute for converting to a SkyCoordinate." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from astropy_xarray.coordinates import dataset_to_skycoord, load_frame\n", + "\n", + "ds = skycoord_to_dataset(sky_position, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])])\n", + "\n", + "# frame specific info stored as dataset attribute\n", + "assert load_frame(ds.attrs[\"frame\"]) == sky_position.replicate_without_data()\n", + "\n", + "# dataset\n", + "result = dataset_to_skycoord(ds)\n", + "np.testing.assert_array_equal(result.ra, sky_position.ra)\n", + "np.testing.assert_array_equal(result.dec, sky_position.dec)\n", + "np.testing.assert_array_equal(result.distance, sky_position.distance)\n", + "\n", + "result" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Serialization and Deserialization\n", + "\n", + "Using dequantifing, all type metadata can be converted entriely to json-like attributes:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# serializable layout\n", + "dds = skycoord_to_dataset(\n", + " sky_direction, coords=[(\"field\", [0, 1, 2, 3])]\n", + ").astropy.dequantify()\n", + "dds" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This allows for compatibility with simple serialization implementations (no custom type handling required)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import msgpack_numpy\n", + "\n", + "\n", + "def to_msgpack_numpy(ds: xr.Dataset):\n", + " return msgpack_numpy.packb(ds.to_dict(data=\"array\"))\n", + "\n", + "\n", + "def from_msgpack_numpy(buf: memoryview):\n", + " return xr.Dataset.from_dict(msgpack_numpy.unpackb(buf))\n", + "\n", + "\n", + "sc = SkyCoord(ra=np.random.random(1000) * u.rad, dec=np.random.random(1000) * u.rad)\n", + "ds = skycoord_to_dataset(sc)\n", + "display(ds)\n", + "\n", + "raw = to_msgpack_numpy(ds.astropy.dequantify())\n", + "display(raw)\n", + "\n", + "result = from_msgpack_numpy(raw).astropy.quantify()\n", + "display(result)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## SkyCoordinate Quirks" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# SkyCoord itself doesn't strongly infer directional reprentation and differential types.\n", + "import astropy.units as u\n", + "from astropy.coordinates import SkyCoord\n", + "\n", + "sky_direction = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + ")\n", + "assert (\n", + " sky_direction.representation_type.name == \"spherical\"\n", + ") # would expect \"unitspherical\"\n", + "assert (\n", + " sky_direction.differential_type.name == \"sphericalcoslat\"\n", + ") # would expect \"unitsphericalcoslat\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# For differential mismatch, this effects SkyCoord methods until the prefered representations are cached.\n", + "sky_position = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " distance=[1, 1, 1, 1] * u.dimensionless_unscaled,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + ")\n", + "assert sky_position.representation_type.name == \"spherical\"\n", + "assert sky_position.differential_type.name == \"sphericalcoslat\"\n", + "sky_position_old_str = str(sky_position)\n", + "\n", + "sky_position.represent_as(\"spherical\", \"sphericalcoslat\") # cache updated\n", + "assert sky_position.representation_type.name == \"spherical\"\n", + "assert sky_position.differential_type.name == \"sphericalcoslat\"\n", + "assert str(sky_position) != sky_position_old_str # string method changed" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# To avoid mutable cache quirks, use explicit repr and diff types for unit sphericals\n", + "sky_position = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " distance=[1, 1, 1, 1] * u.dimensionless_unscaled,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=ICRS(\n", + " representation_type=\"spherical\",\n", + " differential_type=\"unitsphericalcoslat\",\n", + " ),\n", + ")\n", + "sky_position_old_str = str(sky_position)\n", + "sky_position.represent_as(\"spherical\", \"sphericalcoslat\")\n", + "assert str(sky_position) == sky_position_old_str" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sky_direction = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + ")\n", + "\n", + "# NOTE: astropy skycoord frame casts to default non-unit representation and may alias data\n", + "print(\"SkyCoord:\", sky_direction.ra.data[1])\n", + "\n", + "# dataset serialization in FRAME mode avoids unit conversions where possible\n", + "ds = skycoord_to_dataset(sky_direction, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])])\n", + "print(\"Dataset:\", ds.ra.data[1].value)\n", + "\n", + "\n", + "sky_position = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " distance=[1, 1, 1, 1] * u.dimensionless_unscaled,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + ")\n", + "\n", + "display(sky_position)\n", + "ds = skycoord_to_dataset(sky_position, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])])\n", + "display(ds)\n", + "\n", + "# NOTE: positional representation avoids aliasing\n", + "assert sky_position.ra.data[1] == ds.ra.data[1].value" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# TODO" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import astropy.units as u\n", + "import numpy as np\n", + "import xarray as xr\n", + "from astropy.coordinates import ICRS, SkyCoord\n", + "\n", + "from astropy_xarray.coordinates.sky_coord import (\n", + " DatasetRepresentation,\n", + " dataset_to_skycoord,\n", + " skycoord_to_dataset,\n", + ")\n", + "\n", + "sky_direction = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + ")\n", + "assert sky_direction.representation_type.name == \"spherical\"\n", + "assert sky_direction.differential_type.name == \"sphericalcoslat\"\n", + "\n", + "\n", + "# display(sky_direction.distance)\n", + "# display(sky_direction.radial_velocity)\n", + "\n", + "# sky_direction.data.distance\n", + "# sky_direction.data.differentials.get(\"s\").d_distance\n", + "# display(skycoord_to_dataset(sky_direction, coords=[(\"field\", [0, 1, 2, 3])], representation=DatasetRepresentation.FRAME))\n", + "\n", + "# TODO support different dim and coord name\n", + "# ds = skycoord_to_dataset(sky_direction, coords=[(\"field\", [0, 1, 2, 3])], representation=DatasetRepresentation.FRAME)\n", + "ds = skycoord_to_dataset(\n", + " sky_direction,\n", + " coords=[(\"field\", [0, 1, 2, 3])],\n", + " representation=DatasetRepresentation.FRAME_DATA,\n", + ")\n", + "display(ds)\n", + "res = dataset_to_skycoord(ds)\n", + "display(res)\n", + "assert res.representation_type.name == \"spherical\"\n", + "assert res.differential_type.name == \"sphericalcoslat\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sky_direction = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " distance=[2, 3, 4, 5] * u.pc,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + ")\n", + "assert sky_direction.frame.representation_type.name == \"spherical\"\n", + "assert sky_direction.frame.differential_type.name == \"sphericalcoslat\"\n", + "\n", + "sky_direction.frame.representation_type = \"cartesian\"\n", + "display(sky_direction)\n", + "assert sky_direction.frame.representation_type.name == \"cartesian\"\n", + "assert sky_direction.frame.differential_type.name == \"cartesian\"\n", + "\n", + "display(sky_direction.represent_as(\"cartesian\"))\n", + "assert type(sky_direction.frame.data).name == \"spherical\"\n", + "\n", + "assert hasattr(sky_direction, \"x\")\n", + "assert hasattr(sky_direction, \"y\")\n", + "assert hasattr(sky_direction, \"z\")\n", + "\n", + "# desired\n", + "# display(sky_direction.frame.data, sky_direction.frame.data.differentials)\n", + "# display(sky_direction.frame.get_representation_component_names())\n", + "# display(sky_direction.frame.represent_as(sky_direction.representation_type))\n", + "\n", + "# failure\n", + "ds = skycoord_to_dataset(\n", + " sky_direction,\n", + " coords=[(\"field\", [0, 1, 2, 3])],\n", + " representation=DatasetRepresentation.FRAME_DATA,\n", + ")\n", + "display(ds)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ds = skycoord_to_dataset(\n", + " sky_direction,\n", + " coords=[(\"field\", [0, 1, 2, 3])],\n", + " representation=DatasetRepresentation.DATA,\n", + ")\n", + "display(ds)\n", + "\n", + "# # TODO handle mismatch (fallback to data?)\n", + "# ds = skycoord_to_dataset(sky_direction, coords=[(\"field\", [0, 1, 2, 3])], representation=DatasetRepresentation.FRAME_DEFAULT)\n", + "# display(ds)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sky_direction = SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " distance=[2, 3, 4, 5] * u.pc,\n", + " pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n", + " # radial_velocity=[1, 1, 1, 1] * u.mas / u.yr,\n", + " frame=\"icrs\",\n", + ")\n", + "assert sky_direction.representation_type.name == \"spherical\"\n", + "assert sky_direction.differential_type.name == \"sphericalcoslat\"\n", + "\n", + "sky_direction.representation_type = \"cartesian\"\n", + "# sky_direction.differential_type = \"cartesian\"\n", + "display(sky_direction)\n", + "assert sky_direction.differential_type.name == \"cartesian\"\n", + "\n", + "sky_direction.represent_as(\"cartesian\", \"cartesian\")\n", + "# TODO fallback to data for missing\n", + "ds = skycoord_to_dataset(\n", + " sky_direction,\n", + " coords=[(\"field\", [0, 1, 2, 3])],\n", + " representation=DatasetRepresentation.FRAME_DATA,\n", + ")\n", + "display(ds)\n", + "assert ds.attrs[\"frame\"][\"data\"][\"representation_type\"] == \"cartesian\", ds.attrs[\n", + " \"frame\"\n", + "][\"data\"][\"representation_type\"]\n", + "\n", + "ds = skycoord_to_dataset(\n", + " sky_direction,\n", + " coords=[(\"field\", [0, 1, 2, 3])],\n", + " representation=DatasetRepresentation.FRAME_DEFAULT,\n", + ")\n", + "display(ds)\n", + "\n", + "ds = skycoord_to_dataset(\n", + " sky_direction,\n", + " coords=[(\"field\", [0, 1, 2, 3])],\n", + " representation=DatasetRepresentation.DATA,\n", + ")\n", + "display(ds)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "type(sky_direction.frame.data)\n", + "type(sky_direction.frame.data.differentials.get(\"s\"))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/whats-new.rst b/docs/whats-new.rst index 48817523..7ce8622c 100644 --- a/docs/whats-new.rst +++ b/docs/whats-new.rst @@ -2,6 +2,12 @@ What's new ========== + +0.2 (TBD) +--------- + +- Added `coordinates` submodule with `dataset_to_skycoord` and `skycoord_to_dataset` excplit conversion functions for working with :py:class:`astropy.coordinates.SkyCoord`. + 0.1 (17 Jul 2025) ----------------- diff --git a/pyproject.toml b/pyproject.toml index 2b731d76..08c9a738 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,9 +25,10 @@ classifiers = [ ] requires-python = ">=3.10" dependencies = [ - "numpy >= 1.23", - "xarray >= 2022.06.0", - "astropy >= 5.0.0", + "numpy>=1.23,<2.3.0", + "xarray>=2022.06.0,<=2025.4.0", + "astropy>=6.1.0,<7.2.0", + "pandas>=2.3.0,<3.0.0", ] dynamic = ["version"] @@ -38,6 +39,12 @@ Documentation = "https://astropy-xarray.readthedocs.io/en/stable" [project.optional-dependencies] test = [ "pytest>=8.0", + "pytest-cov", + "dask[array]", + "msgpack>=1.1.2", + "msgpack_numpy>=0.4.8", + "cf_xarray>=0.10.6", + "bottleneck", ] docs = [ "sphinx>=8.0", @@ -103,8 +110,7 @@ known-first-party = ["astropy_xarray"] known-third-party = ["xarray"] [tool.ruff.lint.flake8-tidy-imports] -# Disallow all relative imports. -ban-relative-imports = "all" +ban-relative-imports = "all" # Disallow all relative imports. [tool.coverage.run] source = ["astropy_xarray"]