From 9ad1577af190cf11a63355f4ac7f4dc08440d2a6 Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Thu, 18 Sep 2025 00:13:37 +0800 Subject: [PATCH 01/10] Add SkyCoordinates and specialized classes Time, TimeDelta, Latitude, Longitude Signed-off-by: Callan Gray --- .pre-commit-config.yaml | 2 +- astropy_xarray/accessors.py | 57 ++- astropy_xarray/conversion.py | 68 +++- astropy_xarray/coordinates/__init__.py | 12 + astropy_xarray/coordinates/core.py | 26 ++ astropy_xarray/coordinates/frame.py | 304 ++++++++++++++ astropy_xarray/coordinates/representation.py | 81 ++++ astropy_xarray/coordinates/sky_coord.py | 124 ++++++ astropy_xarray/formatting.py | 2 + astropy_xarray/tests/multi_index_utils.py | 58 +++ astropy_xarray/tests/test_accessors_coords.py | 117 ++++++ astropy_xarray/tests/test_accessors_log.py | 104 +++++ .../tests/test_accessors_skycoords.py | 289 +++++++++++++ astropy_xarray/tests/test_accessors_time.py | 118 ++++++ astropy_xarray/tests/test_time.py | 117 ++++++ astropy_xarray/tests/test_visibility.py | 378 ++++++++++++++++++ astropy_xarray/time_compat.py | 55 +++ docs/examples/plotting.ipynb | 119 +++++- docs/examples/requirements.txt | 3 + 19 files changed, 2015 insertions(+), 19 deletions(-) create mode 100644 astropy_xarray/coordinates/__init__.py create mode 100644 astropy_xarray/coordinates/core.py create mode 100644 astropy_xarray/coordinates/frame.py create mode 100644 astropy_xarray/coordinates/representation.py create mode 100644 astropy_xarray/coordinates/sky_coord.py create mode 100644 astropy_xarray/tests/multi_index_utils.py create mode 100644 astropy_xarray/tests/test_accessors_coords.py create mode 100644 astropy_xarray/tests/test_accessors_log.py create mode 100644 astropy_xarray/tests/test_accessors_skycoords.py create mode 100644 astropy_xarray/tests/test_accessors_time.py create mode 100644 astropy_xarray/tests/test_time.py create mode 100644 astropy_xarray/tests/test_visibility.py create mode 100644 astropy_xarray/time_compat.py create mode 100644 docs/examples/requirements.txt 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/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..b7aa2253 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,33 @@ 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 "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 +143,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 +176,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 +459,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..93187cba --- /dev/null +++ b/astropy_xarray/coordinates/frame.py @@ -0,0 +1,304 @@ +from types import NoneType + +import astropy.units as u +import numpy as np +from astropy.coordinates import ( + BaseCoordinateFrame, + BaseRepresentation, + CartesianDifferential, + CartesianRepresentation, + CylindricalDifferential, + CylindricalRepresentation, + EarthLocation, + SphericalCosLatDifferential, + SphericalDifferential, + 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.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.name, + "differential_type": frame.differential_type.name, + } + if frame.has_data: + ser["data"] = {"representation_type": frame.data.name} + if "s" in frame.data.differentials: + ser["data"]["differential_type"] = frame.data.differentials["s"].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, + { + 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, data: dict[str, np.ndarray] +) -> BaseRepresentation: + match representation_type: + case "unitspherical": + RepresentationClass = UnitSphericalRepresentation + case "spherical": + RepresentationClass = SphericalRepresentation + case "cartesian": + RepresentationClass = CartesianRepresentation + case "cylindrical": + RepresentationClass = CylindricalRepresentation + case _: + raise NotImplementedError(representation_type) + match differential_type: + case "unitspherical": + DifferentialClass = UnitSphericalDifferential + case "unitsphericalcoslat": + DifferentialClass = UnitSphericalCosLatDifferential + case "spherical": + DifferentialClass = SphericalDifferential + case "sphericalcoslat": + DifferentialClass = SphericalCosLatDifferential + case "cartesian": + DifferentialClass = CartesianDifferential + case "cylindrical": + DifferentialClass = CylindricalDifferential + case None: + DifferentialClass = None + case _: + raise NotImplementedError(differential_type) + 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, + ) 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..b4d15a70 --- /dev/null +++ b/astropy_xarray/coordinates/sky_coord.py @@ -0,0 +1,124 @@ +import itertools +from collections.abc import Generator + +import astropy.units as u +import xarray as xr +from astropy.coordinates import CartesianRepresentation, SkyCoord + +from astropy_xarray.coordinates.frame import dump_frame, load_frame, load_representation + + +def _skycoord_representation_component_names( + skycoord: SkyCoord, + skip_implicit=True, + use_data_components=True, +) -> tuple[str, ...]: + if use_data_components: + return tuple( + name + for name in skycoord.data.components + if not skip_implicit + or skycoord.representation_type == CartesianRepresentation + or getattr(skycoord.data, name).unit is not u.dimensionless_unscaled + ) + else: + return tuple( + name + for name in skycoord.get_representation_component_names() + if not skip_implicit + or skycoord.representation_type == CartesianRepresentation + or getattr(skycoord, name).unit is not u.dimensionless_unscaled + ) + + +def _skycoord_differential_component_names( + skycoord: SkyCoord, + skip_implicit=True, + use_data_components=True, +) -> tuple[str, ...]: + diff = skycoord.data.differentials.get("s") + if diff is None: + return () + if use_data_components: + return tuple( + name + for name in diff.components + if not skip_implicit + or getattr(diff, name).unit is not u.dimensionless_unscaled + ) + else: + return tuple(skycoord.get_representation_component_names("s")) + + +def _skycoord_component_names( + skycoord: SkyCoord, + skip_implicit=True, + use_data_components=True, +) -> Generator[str, None, None]: + if use_data_components: + yield from itertools.chain( + _skycoord_representation_component_names(skycoord, skip_implicit), + _skycoord_differential_component_names(skycoord, skip_implicit), + ) + else: + yield from itertools.chain( + _skycoord_representation_component_names(skycoord, skip_implicit), + _skycoord_differential_component_names(skycoord, skip_implicit), + ) + + +def _skycoord_components( + skycoord: SkyCoord, use_frame_names=True, skip_dimensionless=True +) -> dict[str, u.Quantity]: + return { + name: ( + getattr(skycoord.data, name) + if hasattr(skycoord.data, name) + else getattr(skycoord.data.differentials["s"], name) + ) + for name in _skycoord_component_names( + skycoord, use_frame_names, skip_dimensionless + ) + } + + +def _skycoord_to_dataarrays( + skycoord: SkyCoord, + coords: list[tuple[str, list]] | None = None, + skip_dimensionless=True, +) -> 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, skip_dimensionless + ).items() + } + + +def skycoord_to_dataset( + skycoord: SkyCoord, + coords: list[tuple[str, list[int]]] | None = None, +) -> xr.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: + dsq = ds.astropy.quantify() + frame = load_frame(ds.attrs["frame"]).realize_frame( + representation_type=ds.attrs["frame"]["representation_type"], + differential_type=ds.attrs["frame"]["differential_type"], + data=load_representation( + ds.attrs["frame"]["data"]["representation_type"], + ds.attrs["frame"]["data"].get("differential_type"), + {k: v.data for k, v in dsq.data_vars.items()}, + ), + ) + 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/multi_index_utils.py b/astropy_xarray/tests/multi_index_utils.py new file mode 100644 index 00000000..449ba347 --- /dev/null +++ b/astropy_xarray/tests/multi_index_utils.py @@ -0,0 +1,58 @@ +import cf_xarray as cfxr +import pandas as pd +import xarray as xr + + +def compress_multindex_in_ds(ds: xr.Dataset) -> xr.Dataset: + multiindex_coords = [ + coord + for coord in ds.coords + if isinstance(ds.coords[coord].to_index(), pd.MultiIndex) + ] + if multiindex_coords: + ds = cfxr.encode_multi_index_as_compress(ds, multiindex_coords) + return ds + + +def reset_multindex_in_tree(tree: xr.DataTree) -> xr.DataTree: + # Create a new tree with MultiIndex reset in each node's Dataset + def process_node(node: xr.DataTree) -> xr.DataTree: + node_data = ( + compress_multindex_in_ds(node.dataset) if node.dataset is not None else None + ) + + new_node = xr.DataTree(dataset=node_data, name=node.name) + for child_name, child in node.children.items(): + new_node[child_name] = process_node(child) + return new_node + + return process_node(tree) + + +def save_datatree_compress_multi_index(tree: xr.DataTree, path: str, **kwargs): + reset_multindex_in_tree(tree).to_zarr(path, **kwargs) + + +def restore_multindex_in_ds(ds: xr.Dataset) -> xr.Dataset: + # Only restore if all coordinate variables are present + if any("compress" in coord.attrs for coord in ds.coords.values()): + ds = cfxr.decode_compress_to_multi_index(ds) + return ds + + +def restore_multindex_in_tree(tree: xr.Dataset) -> xr.Dataset: + # Recursively restore MultiIndexes in each dataset + def process_node(node): + if node.dataset is not None: + node.dataset = restore_multindex_in_ds(node.dataset) + for child in node.children.values(): + process_node(child) + return node + + return process_node(tree) + + +def open_datatree_decompress_multi_index(path: str, **kwargs): + return restore_multindex_in_tree( + xr.open_datatree(path, **kwargs), + ) 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_skycoords.py b/astropy_xarray/tests/test_accessors_skycoords.py new file mode 100644 index 00000000..564170ee --- /dev/null +++ b/astropy_xarray/tests/test_accessors_skycoords.py @@ -0,0 +1,289 @@ +from collections.abc import Generator + +import astropy.units as u +import numpy as np +import pytest +import xarray as xr +import zarr.storage as zs +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 zarr.abc.store import Store + +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, +) + + +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 _skycoord_representation_components(skycoord: SkyCoord) -> tuple[str, ...]: + return tuple(name for name in skycoord._representation) + + +@pytest.mark.parametrize( + ("frame", "expected_base", "expected_s"), + [ + (ICRS(), ("ra", "dec"), ("pm_ra_cosdec", "pm_dec", "radial_velocity")), + (AltAz(), ("az", "alt"), ("pm_az_cosalt", "pm_alt", "radial_velocity")), + (ITRS(), ("x", "y", "z"), ("v_x", "v_y", "v_z")), + ], +) +def test_direction_get_components(frame, expected_base, expected_s): + + rep = UnitSphericalRepresentation( + [[0.1], [2], [0.2]] * u.deg, + [[0.5], [7], [0.7]] * u.deg, + differentials=UnitSphericalDifferential( + d_lon=[[0.002], [0.002], [0.002]] * u.deg / u.yr, + d_lat=[[0.002], [0.002], [0.002]] * u.deg / u.yr, + ), + ) + sc = SkyCoord(frame.realize_frame(rep)) + assert _skycoord_representation_component_names(sc, True, False) == expected_base + assert _skycoord_differential_component_names(sc, True, False) == expected_s + + +@pytest.fixture(name="store", scope="session") +def store_fixture() -> Generator[Store, None, None]: + with zs.MemoryStore() as store: + yield store + + +@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( + store, + frame: BaseCoordinateFrame, + data: BaseRepresentation, + expected_data_classes: tuple[type, ...], +): + # all frames support all representation types + representation = "cylindrical" + + expected = SkyCoord(frame.realize_frame(data)) + expected.representation_type = representation + expected.differential_type = representation + assert expected.frame.data.name == data.name + assert expected.frame.data.differentials["s"].name == data.differentials["s"].name + + coords = [ + ("calibrator_name", ["a1", "a2"]), + ("time_polynomial", [0]), + ] + ds = skycoord_to_dataset(expected, coords=coords) + assert len(ds.data_vars) == ( + len(expected.frame.data.components) + + len(expected.frame.data.differentials["s"].components) + ) + + expected_classes = zip(expected.frame.data.components, expected_data_classes) + for expected_key, expected_class_type in expected_classes: + assert isinstance(ds.data_vars[expected_key].data, expected_class_type) + assert ds.data_vars[expected_key].coords.identical(xr.Coordinates(dict(coords))) + + # sanity checks + assert ds.attrs["frame"]["representation_type"] == representation + assert ds.attrs["frame"]["differential_type"] == representation + assert ds.attrs["frame"]["data"]["representation_type"] == data.name + assert ( + ds.attrs["frame"]["data"]["differential_type"] == data.differentials["s"].name + ) + + actual = dataset_to_skycoord(ds) + + # check representations + assert actual.representation_type.name == representation + assert actual.differential_type.name == representation + assert actual.frame.data.name == data.name + assert actual.frame.data.differentials["s"].name == data.differentials["s"].name + + # 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 + print(repr(actual.frame), "\n\n", repr(expected.frame)) + assert actual.is_equivalent_frame(expected) + assert ds.coords.equals(xr.Coordinates(dict(coords))) + np.testing.assert_array_equal(actual, expected) + + # Check intermediate dataset is serializable + ds.astropy.dequantify().to_zarr(store, mode="w", consolidated=True) + ds_store = xr.open_dataset( + store, engine="zarr", consolidated=True + ).astropy.quantify() + xr.testing.assert_identical(ds, ds_store) 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_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]), + ], + ) + ), + ), + ) + + m = astropy_encode_msgpack(VISIBILITY) + ds_out = astropy_decode_msgpack(m) + xr.testing.assert_identical(VISIBILITY, ds_out) + + store = zs.MemoryStore() + + VISIBILITY.astropy.dequantify().to_zarr(store, mode="w", consolidated=True) + dt_store = xr.open_datatree( + store, engine="zarr", consolidated=True + ).astropy.quantify() + + xr.testing.assert_identical(VISIBILITY, dt_store) + + expected = dataset_to_skycoord(VISIBILITY.field_phase_center.dataset) + actual = dataset_to_skycoord(dt_store.field_phase_center.dataset) + np.testing.assert_array_equal(actual, expected, strict=True) + + +@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"] + ), + **xr.Coordinates.from_pandas_multiindex( + pd.MultiIndex.from_tuples( + generate_baselines(a), names=("antenna1", "antenna2") + ), + dim="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( + # SphericalRepresentation( + # lon=[[0.1, 0.2]] * u.deg, + # lat=[[0.5, 0.2]] * u.deg, + # distance=[[0, 0]] * u.pc, + # ).with_differentials( + # SphericalDifferential( + # d_lon=[[1, 1]] * u.deg / u.s, + # d_lat=[[1, 1]] * u.deg / u.s, + # d_distance=[[1, 1]] * u.pc / u.yr, + # ) + # ), + # frame=ICRS() + # ), + 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", [0]), + ("time_poly", [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", [0]), + ], + ) + ), + ), + ) + + m = astropy_encode_msgpack(VISIBILITY) + ds_out = astropy_decode_msgpack(m) + xr.testing.assert_identical(VISIBILITY, ds_out) + + store = zs.MemoryStore() + + save_datatree_compress_multi_index( + VISIBILITY.astropy.dequantify(), store, mode="w", consolidated=True + ) + dt_store = open_datatree_decompress_multi_index( + store, engine="zarr" + ).astropy.quantify() + + xr.testing.assert_identical(VISIBILITY, dt_store) + + expected = dataset_to_skycoord(VISIBILITY.field_phase_center.dataset) + actual = dataset_to_skycoord(dt_store.field_phase_center.dataset) + + np.testing.assert_array_equal(actual, expected, strict=True) 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/examples/plotting.ipynb b/docs/examples/plotting.ipynb index 84578c3d..4a8f9fee 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,116 @@ "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.sky_coord import dataset_to_skycoord, skycoord_to_dataset\n", + "\n", + "s = skycoord_to_dataset(\n", + " SkyCoord(\n", + " ra=[2, 6, 7, 4] * u.deg,\n", + " dec=[4, 7, 4, 3] * u.deg,\n", + " frame=\"icrs\",\n", + " ),\n", + " coords=[(\"field\", [0, 1, 2, 3])],\n", + ")\n", + "s" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "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 From 6b90258baf0adf183355a463e61575cfb744b420 Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Tue, 18 Nov 2025 13:17:44 +0800 Subject: [PATCH 02/10] all frames tested Signed-off-by: Callan Gray --- astropy_xarray/__init__.py | 1 + astropy_xarray/conversion.py | 2 + astropy_xarray/coordinates/frame.py | 113 ++-- astropy_xarray/coordinates/sky_coord.py | 163 +++--- .../tests/test_accessors_skycoords.py | 289 ----------- astropy_xarray/tests/test_frame_rep.py | 124 +++++ astropy_xarray/tests/test_sky_coord.py | 427 +++++++++++++++ docs/examples/plotting.ipynb | 33 +- docs/examples/skycoord.ipynb | 484 ++++++++++++++++++ 9 files changed, 1224 insertions(+), 412 deletions(-) delete mode 100644 astropy_xarray/tests/test_accessors_skycoords.py create mode 100644 astropy_xarray/tests/test_frame_rep.py create mode 100644 astropy_xarray/tests/test_sky_coord.py create mode 100644 docs/examples/skycoord.ipynb diff --git a/astropy_xarray/__init__.py b/astropy_xarray/__init__.py index 8f17bca7..dc1cc078 100644 --- a/astropy_xarray/__init__.py +++ b/astropy_xarray/__init__.py @@ -17,6 +17,7 @@ import astropy.units +import astropy_xarray.time_compat from astropy_xarray import accessors, formatting, testing # noqa: F401 from astropy_xarray.index import AstropyIndex diff --git a/astropy_xarray/conversion.py b/astropy_xarray/conversion.py index b7aa2253..79538959 100644 --- a/astropy_xarray/conversion.py +++ b/astropy_xarray/conversion.py @@ -89,6 +89,8 @@ def array_attach_unit(data, unit) -> AstropyType: 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": diff --git a/astropy_xarray/coordinates/frame.py b/astropy_xarray/coordinates/frame.py index 93187cba..11341c54 100644 --- a/astropy_xarray/coordinates/frame.py +++ b/astropy_xarray/coordinates/frame.py @@ -7,15 +7,11 @@ BaseRepresentation, CartesianDifferential, CartesianRepresentation, - CylindricalDifferential, CylindricalRepresentation, EarthLocation, - SphericalCosLatDifferential, - SphericalDifferential, SphericalRepresentation, - UnitSphericalCosLatDifferential, - UnitSphericalDifferential, - UnitSphericalRepresentation, + frame_transform_graph, + representation, ) from astropy.coordinates.builtin_frames import ( CIRS, @@ -247,6 +243,7 @@ def load_frame(frame_dict: dict, with_data: bool = False) -> BaseCoordinateFrame data=load_representation( representation_type, differential_type, + None, { k: u.Quantity(**v) for k, v in frame_dict.get("data").items() @@ -258,47 +255,71 @@ def load_frame(frame_dict: dict, with_data: bool = False) -> BaseCoordinateFrame def load_representation( - representation_type: str, differential_type: str | None, data: dict[str, np.ndarray] + representation_type: str, + differential_type: str | None, + frame_name: str | None, + data: dict[str, np.ndarray], ) -> BaseRepresentation: - match representation_type: - case "unitspherical": - RepresentationClass = UnitSphericalRepresentation - case "spherical": - RepresentationClass = SphericalRepresentation - case "cartesian": - RepresentationClass = CartesianRepresentation - case "cylindrical": - RepresentationClass = CylindricalRepresentation - case _: - raise NotImplementedError(representation_type) - match differential_type: - case "unitspherical": - DifferentialClass = UnitSphericalDifferential - case "unitsphericalcoslat": - DifferentialClass = UnitSphericalCosLatDifferential - case "spherical": - DifferentialClass = SphericalDifferential - case "sphericalcoslat": - DifferentialClass = SphericalCosLatDifferential - case "cartesian": - DifferentialClass = CartesianDifferential - case "cylindrical": - DifferentialClass = CylindricalDifferential - case None: - DifferentialClass = None - case _: - raise NotImplementedError(differential_type) - differentials = ( - DifferentialClass( - **{key: data[key] for key in DifferentialClass.attr_classes.keys()}, + 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, ) - if DifferentialClass is not None - else None - ) + 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 {} + ) - return RepresentationClass( - **{key: data[key] for key in RepresentationClass.attr_classes.keys()}, - differentials=differentials, - copy=False, - ) + 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/sky_coord.py b/astropy_xarray/coordinates/sky_coord.py index b4d15a70..185d20fc 100644 --- a/astropy_xarray/coordinates/sky_coord.py +++ b/astropy_xarray/coordinates/sky_coord.py @@ -1,106 +1,126 @@ import itertools from collections.abc import Generator +from enum import Enum, auto import astropy.units as u +import numpy as np import xarray as xr -from astropy.coordinates import CartesianRepresentation, SkyCoord +from astropy.coordinates import SkyCoord from astropy_xarray.coordinates.frame import dump_frame, load_frame, load_representation +class DatasetRepresentation(Enum): + FRAME = auto() + """ + Frame selected repesentation and differential component names, simplifying for any `unitspherical` types in data. + + Same as `get_representation_component_names()` and `get_differential_component_names()` + """ + + FRAME_FULL = auto() + """ + Frame selected repesentation and differential component names. + + Same as `get_representation_component_names()` and `get_differential_component_names()` + """ + + FRAME_DEFAULT = auto() + """Frame default components names.""" + + FRAME_DATA = auto() + """ + Frame component names from resetting the display representation and differential to match the data. + + Guarentees roundtrip accuracy. + """ + + DATA = auto() + """Data representation components. Guarentees roundtrip accuracy.""" + + def _skycoord_representation_component_names( skycoord: SkyCoord, - skip_implicit=True, - use_data_components=True, ) -> tuple[str, ...]: - if use_data_components: - return tuple( - name - for name in skycoord.data.components - if not skip_implicit - or skycoord.representation_type == CartesianRepresentation - or getattr(skycoord.data, name).unit is not u.dimensionless_unscaled - ) - else: - return tuple( - name - for name in skycoord.get_representation_component_names() - if not skip_implicit - or skycoord.representation_type == CartesianRepresentation - or getattr(skycoord, name).unit is not u.dimensionless_unscaled - ) + 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, - skip_implicit=True, - use_data_components=True, ) -> tuple[str, ...]: - diff = skycoord.data.differentials.get("s") - if diff is None: + if skycoord.data.differentials.get("s") is None: return () - if use_data_components: - return tuple( - name - for name in diff.components - if not skip_implicit - or getattr(diff, name).unit is not u.dimensionless_unscaled - ) - else: - return tuple(skycoord.get_representation_component_names("s")) + + 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, - skip_implicit=True, - use_data_components=True, ) -> Generator[str, None, None]: - if use_data_components: - yield from itertools.chain( - _skycoord_representation_component_names(skycoord, skip_implicit), - _skycoord_differential_component_names(skycoord, skip_implicit), - ) - else: - yield from itertools.chain( - _skycoord_representation_component_names(skycoord, skip_implicit), - _skycoord_differential_component_names(skycoord, skip_implicit), - ) + 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 -def _skycoord_components( - skycoord: SkyCoord, use_frame_names=True, skip_dimensionless=True -) -> dict[str, u.Quantity]: + # 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(skycoord.data, name) - if hasattr(skycoord.data, name) - else getattr(skycoord.data.differentials["s"], name) - ) - for name in _skycoord_component_names( - skycoord, use_frame_names, skip_dimensionless + 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, - skip_dimensionless=True, ) -> 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, skip_dimensionless - ).items() + for name, component in _skycoord_components(skycoord).items() } def skycoord_to_dataset( skycoord: SkyCoord, - coords: list[tuple[str, list[int]]] | None = None, + 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), @@ -110,15 +130,24 @@ def skycoord_to_dataset( ) -def dataset_to_skycoord(ds: xr.Dataset) -> SkyCoord: +def dataset_to_skycoord(ds: xr.Dataset, use_frame_names: bool = True) -> 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"]).realize_frame( - representation_type=ds.attrs["frame"]["representation_type"], - differential_type=ds.attrs["frame"]["differential_type"], - data=load_representation( - ds.attrs["frame"]["data"]["representation_type"], - ds.attrs["frame"]["data"].get("differential_type"), - {k: v.data for k, v in dsq.data_vars.items()}, - ), + 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"] if use_frame_names else None, + {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/tests/test_accessors_skycoords.py b/astropy_xarray/tests/test_accessors_skycoords.py deleted file mode 100644 index 564170ee..00000000 --- a/astropy_xarray/tests/test_accessors_skycoords.py +++ /dev/null @@ -1,289 +0,0 @@ -from collections.abc import Generator - -import astropy.units as u -import numpy as np -import pytest -import xarray as xr -import zarr.storage as zs -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 zarr.abc.store import Store - -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, -) - - -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 _skycoord_representation_components(skycoord: SkyCoord) -> tuple[str, ...]: - return tuple(name for name in skycoord._representation) - - -@pytest.mark.parametrize( - ("frame", "expected_base", "expected_s"), - [ - (ICRS(), ("ra", "dec"), ("pm_ra_cosdec", "pm_dec", "radial_velocity")), - (AltAz(), ("az", "alt"), ("pm_az_cosalt", "pm_alt", "radial_velocity")), - (ITRS(), ("x", "y", "z"), ("v_x", "v_y", "v_z")), - ], -) -def test_direction_get_components(frame, expected_base, expected_s): - - rep = UnitSphericalRepresentation( - [[0.1], [2], [0.2]] * u.deg, - [[0.5], [7], [0.7]] * u.deg, - differentials=UnitSphericalDifferential( - d_lon=[[0.002], [0.002], [0.002]] * u.deg / u.yr, - d_lat=[[0.002], [0.002], [0.002]] * u.deg / u.yr, - ), - ) - sc = SkyCoord(frame.realize_frame(rep)) - assert _skycoord_representation_component_names(sc, True, False) == expected_base - assert _skycoord_differential_component_names(sc, True, False) == expected_s - - -@pytest.fixture(name="store", scope="session") -def store_fixture() -> Generator[Store, None, None]: - with zs.MemoryStore() as store: - yield store - - -@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( - store, - frame: BaseCoordinateFrame, - data: BaseRepresentation, - expected_data_classes: tuple[type, ...], -): - # all frames support all representation types - representation = "cylindrical" - - expected = SkyCoord(frame.realize_frame(data)) - expected.representation_type = representation - expected.differential_type = representation - assert expected.frame.data.name == data.name - assert expected.frame.data.differentials["s"].name == data.differentials["s"].name - - coords = [ - ("calibrator_name", ["a1", "a2"]), - ("time_polynomial", [0]), - ] - ds = skycoord_to_dataset(expected, coords=coords) - assert len(ds.data_vars) == ( - len(expected.frame.data.components) - + len(expected.frame.data.differentials["s"].components) - ) - - expected_classes = zip(expected.frame.data.components, expected_data_classes) - for expected_key, expected_class_type in expected_classes: - assert isinstance(ds.data_vars[expected_key].data, expected_class_type) - assert ds.data_vars[expected_key].coords.identical(xr.Coordinates(dict(coords))) - - # sanity checks - assert ds.attrs["frame"]["representation_type"] == representation - assert ds.attrs["frame"]["differential_type"] == representation - assert ds.attrs["frame"]["data"]["representation_type"] == data.name - assert ( - ds.attrs["frame"]["data"]["differential_type"] == data.differentials["s"].name - ) - - actual = dataset_to_skycoord(ds) - - # check representations - assert actual.representation_type.name == representation - assert actual.differential_type.name == representation - assert actual.frame.data.name == data.name - assert actual.frame.data.differentials["s"].name == data.differentials["s"].name - - # 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 - print(repr(actual.frame), "\n\n", repr(expected.frame)) - assert actual.is_equivalent_frame(expected) - assert ds.coords.equals(xr.Coordinates(dict(coords))) - np.testing.assert_array_equal(actual, expected) - - # Check intermediate dataset is serializable - ds.astropy.dequantify().to_zarr(store, mode="w", consolidated=True) - ds_store = xr.open_dataset( - store, engine="zarr", consolidated=True - ).astropy.quantify() - xr.testing.assert_identical(ds, ds_store) diff --git a/astropy_xarray/tests/test_frame_rep.py b/astropy_xarray/tests/test_frame_rep.py new file mode 100644 index 00000000..f312e48d --- /dev/null +++ b/astropy_xarray/tests/test_frame_rep.py @@ -0,0 +1,124 @@ +import astropy.units as u +import numpy as np +import numpy.testing +import pytest +import xarray as xr +import zarr.storage as zs +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 zarr.abc.store import Store + +from astropy_xarray.coordinates import ( + dataset_to_skycoord, + skycoord_to_dataset, +) +from astropy_xarray.coordinates.sky_coord import ( + DatasetRepresentation, + _skycoord_component_names, + _skycoord_differential_component_names, + _skycoord_representation_component_names, +) + + +def test_regression(): + 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_convert_regression(): + 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, + ), + ) + ) + ) + + skycoord_to_dataset(sc, None, representation=DatasetRepresentation.FRAME) + skycoord_to_dataset(sc, None, representation=DatasetRepresentation.FRAME_DEFAULT) + + sc2 = SkyCoord( + ICRS().realize_frame( + sc.represent_as( + sc.frame.default_representation, sc.frame.default_differential + ) + ) + ) + + assert list(_skycoord_component_names(sc2, True, True)) == list( + _skycoord_component_names(sc2, True, False) + ) + + print(skycoord_to_dataset(sc2, None, representation=DatasetRepresentation.FRAME)) + print( + skycoord_to_dataset( + sc2, None, representation=DatasetRepresentation.FRAME_DEFAULT + ) + ) diff --git a/astropy_xarray/tests/test_sky_coord.py b/astropy_xarray/tests/test_sky_coord.py new file mode 100644 index 00000000..ab89afae --- /dev/null +++ b/astropy_xarray/tests/test_sky_coord.py @@ -0,0 +1,427 @@ +from collections.abc import Generator + +import astropy.units as u +import numpy as np +import pytest +import xarray as xr +import zarr.storage as zs +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 zarr.abc.store import Store + +from astropy_xarray.coordinates import ( + dataset_to_skycoord, + skycoord_to_dataset, +) +from astropy_xarray.coordinates.sky_coord import ( + DatasetRepresentation, + _skycoord_differential_component_names, + _skycoord_representation_component_names, +) + +# TODO: Test without diffs! + + +@pytest.mark.parametrize( + ("representation", "expected_base", "expected_s"), + [ + ( + DatasetRepresentation.FRAME_DATA, + ("ra", "dec"), + ("pm_ra_cosdec", "pm_dec", "radial_velocity"), + ), + ], +) +def test_unitspherical_repr_components(representation, 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( + ("representation", "expected_base", "expected_s"), + [ + ( + DatasetRepresentation.FRAME, + ("ra", "dec", "distance"), + ("pm_ra_cosdec", "pm_dec"), + ), + ( + DatasetRepresentation.FRAME_DEFAULT, + ("ra", "dec", "distance"), + ("pm_ra_cosdec", "pm_dec", "radial_velocity"), + ), + ( + DatasetRepresentation.DATA, + ("lon", "lat", "distance"), + ("d_lon_coslat", "d_lat"), + ), + ], +) +def test_unitspherical_diff_components(representation, 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) + _skycoord_representation_component_names(sc) == expected_base + _skycoord_differential_component_names(sc) == expected_s + assert str(sc) == expected_str + + +@pytest.mark.parametrize( + ("representation", "frame", "expected_base", "expected_s"), + [ + ( + DatasetRepresentation.FRAME_DATA, + ICRS(), + ("ra", "dec"), + ("pm_ra_cosdec", "pm_dec"), + ), + ( + DatasetRepresentation.FRAME_DATA, + AltAz(), + ("az", "alt"), + ("pm_az_cosalt", "pm_alt"), + ), + ( + DatasetRepresentation.FRAME_DATA, + ITRS(), + ("lon", "lat"), + ("pm_lon_coslat", "pm_lat"), + ), + ], +) +def test_spherical_direction_components_realize_frame( + representation, 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( + ("representation", "expected_base", "expected_s"), + [ + (DatasetRepresentation.FRAME_DATA, ("x", "y", "z"), ("v_x", "v_y", "v_z")), + ], +) +def test_cartesian_direction_components(representation, 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.fixture(name="store", scope="session") +def store_fixture() -> Generator[Store, None, None]: + with zs.MemoryStore() as store: + yield store + + +@pytest.mark.parametrize( + "representation", + [ + pytest.param(DatasetRepresentation.FRAME_DATA), + pytest.param(DatasetRepresentation.FRAME_DEFAULT), + ], +) +@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( + store, + frame: BaseCoordinateFrame, + data: BaseRepresentation, + expected_data_classes: tuple[type, ...], + representation: DatasetRepresentation, +): + # all frames support all representation types + representation_type = type(data) + differential_type = type(data.differentials["s"]) + + if representation == DatasetRepresentation.FRAME_DATA: + expected = SkyCoord(frame.realize_frame(data)) + elif representation == DatasetRepresentation.FRAME_DEFAULT: + expected = SkyCoord( + frame.realize_frame( + data.represent_as( + frame.default_representation, frame.default_differential + ) + ) + ) + else: + raise NotImplementedError() + + if representation == DatasetRepresentation.FRAME_DEFAULT: + expected_frame_data_rep_name = frame.default_representation.name + expected_frame_data_diff_name = frame.default_differential.name + elif representation == DatasetRepresentation.FRAME_DATA: + expected_frame_data_rep_name = representation_type.name + expected_frame_data_diff_name = differential_type.name + else: + raise NotImplementedError() + + assert expected.frame.data.name == expected_frame_data_rep_name + assert expected.frame.data.differentials["s"].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 + + # data_var value types + if representation in ( + DatasetRepresentation.FRAME, + DatasetRepresentation.FRAME_DEFAULT, + ): + print(expected_keys) + elif representation == DatasetRepresentation.DATA: + expected_classes = zip( + expected.get_representation_component_names(), expected_data_classes + ) + for expected_key, expected_class_type in expected_classes: + assert isinstance(ds.data_vars[expected_key].data, expected_class_type) + assert ds.data_vars[expected_key].coords.identical( + xr.Coordinates(dict(coords)) + ) + + # sanity check dataset + assert ds.attrs["frame"]["representation_type"] == representation_type.name + assert ds.attrs["frame"]["differential_type"] == differential_type.name + assert ds.attrs["frame"]["data"]["representation_type"] == data.name + assert ( + ds.attrs["frame"]["data"]["differential_type"] + == data.differentials["s"].name + ) + + # Convert back + actual = dataset_to_skycoord(ds) + + # check representations + if representation in ( + DatasetRepresentation.FRAME, + DatasetRepresentation.FRAME_DEFAULT, + ): + print(actual) + if representation == DatasetRepresentation.DATA: + assert actual.representation_type.name == representation_type.name + assert actual.differential_type.name == differential_type.name + assert actual.frame.data.name == data.name + assert actual.frame.data.differentials["s"].name == data.differentials["s"].name + + # 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) + + # Check intermediate dataset is serializable + # ds.astropy.dequantify().to_zarr(store, mode="w", consolidated=True) + # ds_store = xr.open_dataset( + # store, engine="zarr", consolidated=True + # ).astropy.quantify() + # xr.testing.assert_identical(ds, ds_store) diff --git a/docs/examples/plotting.ipynb b/docs/examples/plotting.ipynb index 4a8f9fee..b644025c 100644 --- a/docs/examples/plotting.ipynb +++ b/docs/examples/plotting.ipynb @@ -226,17 +226,29 @@ "source": [ "from astropy.coordinates import SkyCoord\n", "\n", - "from astropy_xarray.sky_coord import dataset_to_skycoord, skycoord_to_dataset\n", + "from astropy_xarray.coordinates.sky_coord import (\n", + " dataset_to_skycoord,\n", + " skycoord_to_dataset,\n", + ")\n", "\n", - "s = skycoord_to_dataset(\n", - " SkyCoord(\n", - " ra=[2, 6, 7, 4] * u.deg,\n", - " dec=[4, 7, 4, 3] * u.deg,\n", - " frame=\"icrs\",\n", - " ),\n", - " coords=[(\"field\", [0, 1, 2, 3])],\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", + "s1 = skycoord_to_dataset(sc, coords=[(\"field\", [0, 1, 2, 3])], use_frame_names=False)\n", + "s2 = skycoord_to_dataset(\n", + " sc, coords=[(\"field\", [0, 1, 2, 3])], use_frame_names=True, skip_implicit=False\n", ")\n", - "s" + "s = skycoord_to_dataset(sc, coords=[(\"field\", [0, 1, 2, 3])])\n", + "\n", + "# display(s1)\n", + "# display(s2)\n", + "s.astropy.dequantify()" ] }, { @@ -246,7 +258,8 @@ "metadata": {}, "outputs": [], "source": [ - "dataset_to_skycoord(s)" + "# display(dataset_to_skycoord(s1, use_frame_names=False))\n", + "display(dataset_to_skycoord(s))" ] }, { diff --git a/docs/examples/skycoord.ipynb b/docs/examples/skycoord.ipynb new file mode 100644 index 00000000..1b0c0411 --- /dev/null +++ b/docs/examples/skycoord.ipynb @@ -0,0 +1,484 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## SkyCoord, Frame and Angle as 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 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", + " DatasetRepresentation,\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: for correctness 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 display and data are also detected." + ] + }, + { + "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", + "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 ICRS, 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(\n", + " sky_direction, DatasetRepresentation.FRAME, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])]\n", + ")\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 +} From 1d8f4860cdbc8b59bb745e71ed0a70ac5f613555 Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Tue, 17 Feb 2026 11:06:19 +0800 Subject: [PATCH 03/10] mark dependency versions, clearer test names Signed-off-by: Callan Gray --- astropy_xarray/tests/test_frame_rep.py | 81 +++++++------------------- pyproject.toml | 10 +++- 2 files changed, 28 insertions(+), 63 deletions(-) diff --git a/astropy_xarray/tests/test_frame_rep.py b/astropy_xarray/tests/test_frame_rep.py index f312e48d..77945eaf 100644 --- a/astropy_xarray/tests/test_frame_rep.py +++ b/astropy_xarray/tests/test_frame_rep.py @@ -1,66 +1,21 @@ import astropy.units as u -import numpy as np import numpy.testing -import pytest -import xarray as xr -import zarr.storage as zs 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 zarr.abc.store import Store +from astropy.coordinates.builtin_frames import ICRS from astropy_xarray.coordinates import ( - dataset_to_skycoord, skycoord_to_dataset, ) from astropy_xarray.coordinates.sky_coord import ( - DatasetRepresentation, _skycoord_component_names, - _skycoord_differential_component_names, - _skycoord_representation_component_names, ) -def test_regression(): +def test_frame_conversion_aliasing(): sc = SkyCoord( ICRS().realize_frame( CartesianRepresentation( @@ -85,7 +40,7 @@ def test_regression(): numpy.testing.assert_array_equal(sc.pm_ra_cosdec, sc2.pm_ra_cosdec) -def test_convert_regression(): +def test_frame_conversion_components(): sc = SkyCoord( ICRS().realize_frame( CartesianRepresentation( @@ -101,8 +56,16 @@ def test_convert_regression(): ) ) - skycoord_to_dataset(sc, None, representation=DatasetRepresentation.FRAME) - skycoord_to_dataset(sc, None, representation=DatasetRepresentation.FRAME_DEFAULT) + 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( @@ -112,13 +75,11 @@ def test_convert_regression(): ) ) - assert list(_skycoord_component_names(sc2, True, True)) == list( - _skycoord_component_names(sc2, True, False) - ) - - print(skycoord_to_dataset(sc2, None, representation=DatasetRepresentation.FRAME)) - print( - skycoord_to_dataset( - sc2, None, representation=DatasetRepresentation.FRAME_DEFAULT - ) - ) + assert list(_skycoord_component_names(sc2)) == [ + "ra", + "dec", + "distance", + "pm_ra_cosdec", + "pm_dec", + "radial_velocity", + ] diff --git a/pyproject.toml b/pyproject.toml index 2b731d76..a843e97a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,9 +25,9 @@ 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>=5.0.0,<7.2.0", ] dynamic = ["version"] @@ -38,6 +38,10 @@ Documentation = "https://astropy-xarray.readthedocs.io/en/stable" [project.optional-dependencies] test = [ "pytest>=8.0", + "zarr>=3.0.0,<3.1.0", + "msgpack>=1.1.2", + "msgpack_numpy>=0.4.8", + "cf_xarray>=0.10.11", ] docs = [ "sphinx>=8.0", From 4d6671bec64881163596876a2edee9138ef9327e Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Tue, 17 Feb 2026 11:21:01 +0800 Subject: [PATCH 04/10] use project dependency versions in ci Signed-off-by: Callan Gray --- .github/workflows/ci.yml | 2 +- pyproject.toml | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) 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/pyproject.toml b/pyproject.toml index a843e97a..19d0bba0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,10 +38,13 @@ Documentation = "https://astropy-xarray.readthedocs.io/en/stable" [project.optional-dependencies] test = [ "pytest>=8.0", - "zarr>=3.0.0,<3.1.0", + "pytest-cov", + "dask[array]", + "zarr>=2.0.0,<3.1.0", "msgpack>=1.1.2", "msgpack_numpy>=0.4.8", "cf_xarray>=0.10.11", + "bottleneck", ] docs = [ "sphinx>=8.0", From 7728c20a6281141bda19707614d5a85957bf9314 Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Tue, 17 Feb 2026 11:48:59 +0800 Subject: [PATCH 05/10] astropy 6 skycoords not supported Signed-off-by: Callan Gray --- .gitignore | 2 +- astropy_xarray/tests/test_frame_rep.py | 2 ++ astropy_xarray/tests/test_sky_coord.py | 9 ++++++--- astropy_xarray/tests/test_visibility.py | 8 +++++++- astropy_xarray/tests/utils.py | 1 + pyproject.toml | 4 ++-- 6 files changed, 19 insertions(+), 7 deletions(-) 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/astropy_xarray/tests/test_frame_rep.py b/astropy_xarray/tests/test_frame_rep.py index 77945eaf..92cfc4b5 100644 --- a/astropy_xarray/tests/test_frame_rep.py +++ b/astropy_xarray/tests/test_frame_rep.py @@ -1,5 +1,6 @@ import astropy.units as u import numpy.testing +import pytest from astropy.coordinates import ( CartesianDifferential, CartesianRepresentation, @@ -40,6 +41,7 @@ def test_frame_conversion_aliasing(): numpy.testing.assert_array_equal(sc.pm_ra_cosdec, sc2.pm_ra_cosdec) +@pytest.mark.xfail(reason="astropy 6.0.0 skycoords not supported") def test_frame_conversion_components(): sc = SkyCoord( ICRS().realize_frame( diff --git a/astropy_xarray/tests/test_sky_coord.py b/astropy_xarray/tests/test_sky_coord.py index ab89afae..1fe268ca 100644 --- a/astropy_xarray/tests/test_sky_coord.py +++ b/astropy_xarray/tests/test_sky_coord.py @@ -1,10 +1,10 @@ from collections.abc import Generator +from typing import Any import astropy.units as u import numpy as np import pytest import xarray as xr -import zarr.storage as zs from astropy.coordinates import ( Angle, BaseCoordinateFrame, @@ -47,7 +47,7 @@ Supergalactic, ) from astropy.time import Time -from zarr.abc.store import Store +from xarray.tests import requires_zarr from astropy_xarray.coordinates import ( dataset_to_skycoord, @@ -180,11 +180,14 @@ def test_cartesian_direction_components(representation, expected_base, expected_ @pytest.fixture(name="store", scope="session") -def store_fixture() -> Generator[Store, None, None]: +def store_fixture() -> Generator[Any, None, None]: + import zarr.storage as zs + with zs.MemoryStore() as store: yield store +@requires_zarr @pytest.mark.parametrize( "representation", [ diff --git a/astropy_xarray/tests/test_visibility.py b/astropy_xarray/tests/test_visibility.py index 9223ecf7..2c90de07 100644 --- a/astropy_xarray/tests/test_visibility.py +++ b/astropy_xarray/tests/test_visibility.py @@ -7,13 +7,13 @@ import pandas as pd import pytest import xarray as xr -import zarr.storage as zs from astropy.coordinates import ( GCRS, ICRS, AltAz, SkyCoord, ) +from xarray.tests import requires_zarr from astropy_xarray.coordinates import ( dataset_to_skycoord, @@ -132,6 +132,7 @@ def generate_baselines(antenna_count) -> list[tuple[int, int]]: return baselines +@requires_zarr def test_simple_visibility_dataset(): a = 4 # antennas b = int(a * (a - 1) / 2) # baselines @@ -221,6 +222,8 @@ def test_simple_visibility_dataset(): ds_out = astropy_decode_msgpack(m) xr.testing.assert_identical(VISIBILITY, ds_out) + import zarr.storage as zs + store = zs.MemoryStore() VISIBILITY.astropy.dequantify().to_zarr(store, mode="w", consolidated=True) @@ -235,6 +238,7 @@ def test_simple_visibility_dataset(): np.testing.assert_array_equal(actual, expected, strict=True) +@requires_zarr @pytest.mark.parametrize( "skycoords,frame,unit", [ @@ -361,6 +365,8 @@ def test_visibility_dataset(skycoords: list[SkyCoord], frame, unit): ds_out = astropy_decode_msgpack(m) xr.testing.assert_identical(VISIBILITY, ds_out) + import zarr.storage as zs + store = zs.MemoryStore() save_datatree_compress_multi_index( diff --git a/astropy_xarray/tests/utils.py b/astropy_xarray/tests/utils.py index 08e306e2..64d7e496 100644 --- a/astropy_xarray/tests/utils.py +++ b/astropy_xarray/tests/utils.py @@ -30,6 +30,7 @@ def importorskip(name): has_dask_array, requires_dask_array = importorskip("dask.array") has_scipy, requires_scipy = importorskip("scipy") has_bottleneck, requires_bottleneck = importorskip("bottleneck") +has_zarr, requires_zarr = importorskip("zarr") @contextmanager diff --git a/pyproject.toml b/pyproject.toml index 19d0bba0..ab08ec2e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,10 +40,10 @@ test = [ "pytest>=8.0", "pytest-cov", "dask[array]", - "zarr>=2.0.0,<3.1.0", + "zarr>=3.0.0; python_version >= '3.11'", "msgpack>=1.1.2", "msgpack_numpy>=0.4.8", - "cf_xarray>=0.10.11", + "cf_xarray>=0.10.6", "bottleneck", ] docs = [ From 0a57e8ecaae1b0774b57d4b72640ab97c3707623 Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Wed, 25 Mar 2026 21:32:24 +0800 Subject: [PATCH 06/10] set supported versions Signed-off-by: Callan Gray --- astropy_xarray/tests/test_sky_coord.py | 11 ++-- ...ty.py => test_visibility_serialization.py} | 51 ++++++++++--------- pyproject.toml | 6 +-- 3 files changed, 34 insertions(+), 34 deletions(-) rename astropy_xarray/tests/{test_visibility.py => test_visibility_serialization.py} (90%) diff --git a/astropy_xarray/tests/test_sky_coord.py b/astropy_xarray/tests/test_sky_coord.py index 1fe268ca..8368f54d 100644 --- a/astropy_xarray/tests/test_sky_coord.py +++ b/astropy_xarray/tests/test_sky_coord.py @@ -59,20 +59,17 @@ _skycoord_representation_component_names, ) -# TODO: Test without diffs! - @pytest.mark.parametrize( - ("representation", "expected_base", "expected_s"), + ("expected_base", "expected_s"), [ ( - DatasetRepresentation.FRAME_DATA, ("ra", "dec"), ("pm_ra_cosdec", "pm_dec", "radial_velocity"), ), ], ) -def test_unitspherical_repr_components(representation, expected_base, expected_s): +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, @@ -115,8 +112,8 @@ def test_unitspherical_diff_components(representation, expected_base, expected_s pm_dec=[[0.002], [0.002], [0.002]] * u.deg / u.yr, ) expected_str = str(sc) - _skycoord_representation_component_names(sc) == expected_base - _skycoord_differential_component_names(sc) == expected_s + assert _skycoord_representation_component_names(sc) == expected_base + assert _skycoord_differential_component_names(sc) == expected_s assert str(sc) == expected_str diff --git a/astropy_xarray/tests/test_visibility.py b/astropy_xarray/tests/test_visibility_serialization.py similarity index 90% rename from astropy_xarray/tests/test_visibility.py rename to astropy_xarray/tests/test_visibility_serialization.py index 2c90de07..7b528721 100644 --- a/astropy_xarray/tests/test_visibility.py +++ b/astropy_xarray/tests/test_visibility_serialization.py @@ -222,20 +222,20 @@ def test_simple_visibility_dataset(): ds_out = astropy_decode_msgpack(m) xr.testing.assert_identical(VISIBILITY, ds_out) - import zarr.storage as zs + # import zarr.storage as zs - store = zs.MemoryStore() + # store = zs.MemoryStore() - VISIBILITY.astropy.dequantify().to_zarr(store, mode="w", consolidated=True) - dt_store = xr.open_datatree( - store, engine="zarr", consolidated=True - ).astropy.quantify() + # VISIBILITY.astropy.dequantify().to_zarr(store, mode="w", consolidated=True) + # dt_store = xr.open_datatree( + # store, engine="zarr", consolidated=True + # ).astropy.quantify() - xr.testing.assert_identical(VISIBILITY, dt_store) + # xr.testing.assert_identical(VISIBILITY, dt_store) - expected = dataset_to_skycoord(VISIBILITY.field_phase_center.dataset) - actual = dataset_to_skycoord(dt_store.field_phase_center.dataset) - np.testing.assert_array_equal(actual, expected, strict=True) + # expected = dataset_to_skycoord(VISIBILITY.field_phase_center.dataset) + # actual = dataset_to_skycoord(dt_store.field_phase_center.dataset) + # np.testing.assert_array_equal(actual, expected, strict=True) @requires_zarr @@ -276,13 +276,16 @@ def test_visibility_dataset(skycoords: list[SkyCoord], frame, unit): 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", - ), - uvw_label=xr.DataArray(["u", "v", "w"], dims=["uvw_label"]), + ) ), data_vars=dict( uvw=xr.DataArray( @@ -365,20 +368,20 @@ def test_visibility_dataset(skycoords: list[SkyCoord], frame, unit): ds_out = astropy_decode_msgpack(m) xr.testing.assert_identical(VISIBILITY, ds_out) - import zarr.storage as zs + # import zarr.storage as zs - store = zs.MemoryStore() + # store = zs.MemoryStore() - save_datatree_compress_multi_index( - VISIBILITY.astropy.dequantify(), store, mode="w", consolidated=True - ) - dt_store = open_datatree_decompress_multi_index( - store, engine="zarr" - ).astropy.quantify() + # save_datatree_compress_multi_index( + # VISIBILITY.astropy.dequantify(), store, mode="w", consolidated=True + # ) + # dt_store = open_datatree_decompress_multi_index( + # store, engine="zarr" + # ).astropy.quantify() - xr.testing.assert_identical(VISIBILITY, dt_store) + # xr.testing.assert_identical(VISIBILITY, dt_store) - expected = dataset_to_skycoord(VISIBILITY.field_phase_center.dataset) - actual = dataset_to_skycoord(dt_store.field_phase_center.dataset) + # expected = dataset_to_skycoord(VISIBILITY.field_phase_center.dataset) + # actual = dataset_to_skycoord(dt_store.field_phase_center.dataset) - np.testing.assert_array_equal(actual, expected, strict=True) + # np.testing.assert_array_equal(actual, expected, strict=True) diff --git a/pyproject.toml b/pyproject.toml index ab08ec2e..83815a00 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,16 +18,16 @@ classifiers = [ "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Topic :: Scientific/Engineering", ] -requires-python = ">=3.10" +requires-python = ">=3.11" dependencies = [ "numpy>=1.23,<2.3.0", "xarray>=2022.06.0,<=2025.4.0", - "astropy>=5.0.0,<7.2.0", + "astropy>=7.0.0,<7.2.0", + "pandas>=2.3.0,<3.0.0", ] dynamic = ["version"] From 0c30cfdfb820afa12bd721e9bb8bcce9947f0bc3 Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Wed, 25 Mar 2026 22:50:40 +0800 Subject: [PATCH 07/10] cleanup Signed-off-by: Callan Gray --- astropy_xarray/__init__.py | 2 +- astropy_xarray/coordinates/sky_coord.py | 30 ----- astropy_xarray/tests/multi_index_utils.py | 58 --------- ..._frame_rep.py => test_frame_conversion.py} | 8 +- astropy_xarray/tests/test_sky_coord.py | 117 ++++-------------- .../tests/test_visibility_serialization.py | 83 +++---------- astropy_xarray/tests/utils.py | 1 - docs/examples/skycoord.ipynb | 34 +++-- pyproject.toml | 9 +- 9 files changed, 77 insertions(+), 265 deletions(-) delete mode 100644 astropy_xarray/tests/multi_index_utils.py rename astropy_xarray/tests/{test_frame_rep.py => test_frame_conversion.py} (90%) diff --git a/astropy_xarray/__init__.py b/astropy_xarray/__init__.py index dc1cc078..13e6702c 100644 --- a/astropy_xarray/__init__.py +++ b/astropy_xarray/__init__.py @@ -17,7 +17,7 @@ import astropy.units -import astropy_xarray.time_compat +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/coordinates/sky_coord.py b/astropy_xarray/coordinates/sky_coord.py index 185d20fc..b8c8e30a 100644 --- a/astropy_xarray/coordinates/sky_coord.py +++ b/astropy_xarray/coordinates/sky_coord.py @@ -1,6 +1,5 @@ import itertools from collections.abc import Generator -from enum import Enum, auto import astropy.units as u import numpy as np @@ -10,35 +9,6 @@ from astropy_xarray.coordinates.frame import dump_frame, load_frame, load_representation -class DatasetRepresentation(Enum): - FRAME = auto() - """ - Frame selected repesentation and differential component names, simplifying for any `unitspherical` types in data. - - Same as `get_representation_component_names()` and `get_differential_component_names()` - """ - - FRAME_FULL = auto() - """ - Frame selected repesentation and differential component names. - - Same as `get_representation_component_names()` and `get_differential_component_names()` - """ - - FRAME_DEFAULT = auto() - """Frame default components names.""" - - FRAME_DATA = auto() - """ - Frame component names from resetting the display representation and differential to match the data. - - Guarentees roundtrip accuracy. - """ - - DATA = auto() - """Data representation components. Guarentees roundtrip accuracy.""" - - def _skycoord_representation_component_names( skycoord: SkyCoord, ) -> tuple[str, ...]: diff --git a/astropy_xarray/tests/multi_index_utils.py b/astropy_xarray/tests/multi_index_utils.py deleted file mode 100644 index 449ba347..00000000 --- a/astropy_xarray/tests/multi_index_utils.py +++ /dev/null @@ -1,58 +0,0 @@ -import cf_xarray as cfxr -import pandas as pd -import xarray as xr - - -def compress_multindex_in_ds(ds: xr.Dataset) -> xr.Dataset: - multiindex_coords = [ - coord - for coord in ds.coords - if isinstance(ds.coords[coord].to_index(), pd.MultiIndex) - ] - if multiindex_coords: - ds = cfxr.encode_multi_index_as_compress(ds, multiindex_coords) - return ds - - -def reset_multindex_in_tree(tree: xr.DataTree) -> xr.DataTree: - # Create a new tree with MultiIndex reset in each node's Dataset - def process_node(node: xr.DataTree) -> xr.DataTree: - node_data = ( - compress_multindex_in_ds(node.dataset) if node.dataset is not None else None - ) - - new_node = xr.DataTree(dataset=node_data, name=node.name) - for child_name, child in node.children.items(): - new_node[child_name] = process_node(child) - return new_node - - return process_node(tree) - - -def save_datatree_compress_multi_index(tree: xr.DataTree, path: str, **kwargs): - reset_multindex_in_tree(tree).to_zarr(path, **kwargs) - - -def restore_multindex_in_ds(ds: xr.Dataset) -> xr.Dataset: - # Only restore if all coordinate variables are present - if any("compress" in coord.attrs for coord in ds.coords.values()): - ds = cfxr.decode_compress_to_multi_index(ds) - return ds - - -def restore_multindex_in_tree(tree: xr.Dataset) -> xr.Dataset: - # Recursively restore MultiIndexes in each dataset - def process_node(node): - if node.dataset is not None: - node.dataset = restore_multindex_in_ds(node.dataset) - for child in node.children.values(): - process_node(child) - return node - - return process_node(tree) - - -def open_datatree_decompress_multi_index(path: str, **kwargs): - return restore_multindex_in_tree( - xr.open_datatree(path, **kwargs), - ) diff --git a/astropy_xarray/tests/test_frame_rep.py b/astropy_xarray/tests/test_frame_conversion.py similarity index 90% rename from astropy_xarray/tests/test_frame_rep.py rename to astropy_xarray/tests/test_frame_conversion.py index 92cfc4b5..10c13ff9 100644 --- a/astropy_xarray/tests/test_frame_rep.py +++ b/astropy_xarray/tests/test_frame_conversion.py @@ -1,3 +1,5 @@ +from importlib.metadata import version + import astropy.units as u import numpy.testing import pytest @@ -7,6 +9,7 @@ SkyCoord, ) from astropy.coordinates.builtin_frames import ICRS +from packaging.version import Version from astropy_xarray.coordinates import ( skycoord_to_dataset, @@ -41,7 +44,10 @@ def test_frame_conversion_aliasing(): numpy.testing.assert_array_equal(sc.pm_ra_cosdec, sc2.pm_ra_cosdec) -@pytest.mark.xfail(reason="astropy 6.0.0 skycoords not supported") +@pytest.mark.xfail( + condition=Version(version("astropy")) <= Version("7.0.0"), + reason="astropy 6.0.0 skycoords not supported", +) def test_frame_conversion_components(): sc = SkyCoord( ICRS().realize_frame( diff --git a/astropy_xarray/tests/test_sky_coord.py b/astropy_xarray/tests/test_sky_coord.py index 8368f54d..5713d0a3 100644 --- a/astropy_xarray/tests/test_sky_coord.py +++ b/astropy_xarray/tests/test_sky_coord.py @@ -1,5 +1,4 @@ -from collections.abc import Generator -from typing import Any +import sys import astropy.units as u import numpy as np @@ -47,18 +46,21 @@ Supergalactic, ) from astropy.time import Time -from xarray.tests import requires_zarr from astropy_xarray.coordinates import ( dataset_to_skycoord, skycoord_to_dataset, ) from astropy_xarray.coordinates.sky_coord import ( - DatasetRepresentation, _skycoord_differential_component_names, _skycoord_representation_component_names, ) +pytestmark = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="python3.10 or higher required for skycoords support", +) + @pytest.mark.parametrize( ("expected_base", "expected_s"), @@ -84,26 +86,15 @@ def test_unitspherical_repr_components(expected_base, expected_s): @pytest.mark.parametrize( - ("representation", "expected_base", "expected_s"), + ("expected_base", "expected_s"), [ ( - DatasetRepresentation.FRAME, ("ra", "dec", "distance"), ("pm_ra_cosdec", "pm_dec"), - ), - ( - DatasetRepresentation.FRAME_DEFAULT, - ("ra", "dec", "distance"), - ("pm_ra_cosdec", "pm_dec", "radial_velocity"), - ), - ( - DatasetRepresentation.DATA, - ("lon", "lat", "distance"), - ("d_lon_coslat", "d_lat"), - ), + ) ], ) -def test_unitspherical_diff_components(representation, expected_base, expected_s): +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, @@ -118,31 +109,26 @@ def test_unitspherical_diff_components(representation, expected_base, expected_s @pytest.mark.parametrize( - ("representation", "frame", "expected_base", "expected_s"), + ("frame", "expected_base", "expected_s"), [ ( - DatasetRepresentation.FRAME_DATA, ICRS(), ("ra", "dec"), ("pm_ra_cosdec", "pm_dec"), ), ( - DatasetRepresentation.FRAME_DATA, AltAz(), ("az", "alt"), ("pm_az_cosalt", "pm_alt"), ), ( - DatasetRepresentation.FRAME_DATA, ITRS(), ("lon", "lat"), ("pm_lon_coslat", "pm_lat"), ), ], ) -def test_spherical_direction_components_realize_frame( - representation, frame, expected_base, expected_s -): +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, @@ -157,12 +143,12 @@ def test_spherical_direction_components_realize_frame( @pytest.mark.parametrize( - ("representation", "expected_base", "expected_s"), + ("expected_base", "expected_s"), [ - (DatasetRepresentation.FRAME_DATA, ("x", "y", "z"), ("v_x", "v_y", "v_z")), + (("x", "y", "z"), ("v_x", "v_y", "v_z")), ], ) -def test_cartesian_direction_components(representation, expected_base, expected_s): +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, @@ -176,20 +162,11 @@ def test_cartesian_direction_components(representation, expected_base, expected_ assert _skycoord_differential_component_names(sc) == expected_s -@pytest.fixture(name="store", scope="session") -def store_fixture() -> Generator[Any, None, None]: - import zarr.storage as zs - - with zs.MemoryStore() as store: - yield store - - -@requires_zarr @pytest.mark.parametrize( - "representation", + "use_default_representation", [ - pytest.param(DatasetRepresentation.FRAME_DATA), - pytest.param(DatasetRepresentation.FRAME_DEFAULT), + pytest.param(True), + pytest.param(False), ], ) @pytest.mark.parametrize( @@ -308,19 +285,16 @@ def store_fixture() -> Generator[Any, None, None]: ids=lambda param: param.name, ) def test_skycoord_roundtrip( - store, frame: BaseCoordinateFrame, data: BaseRepresentation, expected_data_classes: tuple[type, ...], - representation: DatasetRepresentation, + use_default_representation: bool, ): # all frames support all representation types representation_type = type(data) differential_type = type(data.differentials["s"]) - if representation == DatasetRepresentation.FRAME_DATA: - expected = SkyCoord(frame.realize_frame(data)) - elif representation == DatasetRepresentation.FRAME_DEFAULT: + if use_default_representation: expected = SkyCoord( frame.realize_frame( data.represent_as( @@ -328,17 +302,12 @@ def test_skycoord_roundtrip( ) ) ) - else: - raise NotImplementedError() - - if representation == DatasetRepresentation.FRAME_DEFAULT: expected_frame_data_rep_name = frame.default_representation.name expected_frame_data_diff_name = frame.default_differential.name - elif representation == DatasetRepresentation.FRAME_DATA: + else: + expected = SkyCoord(frame.realize_frame(data)) expected_frame_data_rep_name = representation_type.name expected_frame_data_diff_name = differential_type.name - else: - raise NotImplementedError() assert expected.frame.data.name == expected_frame_data_rep_name assert expected.frame.data.differentials["s"].name == expected_frame_data_diff_name @@ -360,46 +329,9 @@ def test_skycoord_roundtrip( ) assert set(ds.data_vars) == expected_keys - # data_var value types - if representation in ( - DatasetRepresentation.FRAME, - DatasetRepresentation.FRAME_DEFAULT, - ): - print(expected_keys) - elif representation == DatasetRepresentation.DATA: - expected_classes = zip( - expected.get_representation_component_names(), expected_data_classes - ) - for expected_key, expected_class_type in expected_classes: - assert isinstance(ds.data_vars[expected_key].data, expected_class_type) - assert ds.data_vars[expected_key].coords.identical( - xr.Coordinates(dict(coords)) - ) - - # sanity check dataset - assert ds.attrs["frame"]["representation_type"] == representation_type.name - assert ds.attrs["frame"]["differential_type"] == differential_type.name - assert ds.attrs["frame"]["data"]["representation_type"] == data.name - assert ( - ds.attrs["frame"]["data"]["differential_type"] - == data.differentials["s"].name - ) - # Convert back actual = dataset_to_skycoord(ds) - # check representations - if representation in ( - DatasetRepresentation.FRAME, - DatasetRepresentation.FRAME_DEFAULT, - ): - print(actual) - if representation == DatasetRepresentation.DATA: - assert actual.representation_type.name == representation_type.name - assert actual.differential_type.name == differential_type.name - assert actual.frame.data.name == data.name - assert actual.frame.data.differentials["s"].name == data.differentials["s"].name - # Check Representation for component_name in actual.frame.representation_component_names: np.testing.assert_array_equal( @@ -418,10 +350,3 @@ def test_skycoord_roundtrip( assert actual.is_equivalent_frame(expected) assert ds.coords.equals(xr.Coordinates(dict(coords))) np.testing.assert_array_equal(actual, expected) - - # Check intermediate dataset is serializable - # ds.astropy.dequantify().to_zarr(store, mode="w", consolidated=True) - # ds_store = xr.open_dataset( - # store, engine="zarr", consolidated=True - # ).astropy.quantify() - # xr.testing.assert_identical(ds, ds_store) diff --git a/astropy_xarray/tests/test_visibility_serialization.py b/astropy_xarray/tests/test_visibility_serialization.py index 7b528721..3eef7035 100644 --- a/astropy_xarray/tests/test_visibility_serialization.py +++ b/astropy_xarray/tests/test_visibility_serialization.py @@ -1,4 +1,5 @@ import json +import sys import astropy.units as u import msgpack @@ -13,17 +14,18 @@ AltAz, SkyCoord, ) -from xarray.tests import requires_zarr from astropy_xarray.coordinates import ( dataset_to_skycoord, skycoord_to_dataset, ) -from astropy_xarray.tests.multi_index_utils import ( - open_datatree_decompress_multi_index, - save_datatree_compress_multi_index, + +pytestmark = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="python3.10 or higher required for skycoords support", ) + FIELD_PHASE_CENTER_ICRS = [ SkyCoord( ra=[2, 0.2, 0.02, 0.002] * u.deg, @@ -132,7 +134,6 @@ def generate_baselines(antenna_count) -> list[tuple[int, int]]: return baselines -@requires_zarr def test_simple_visibility_dataset(): a = 4 # antennas b = int(a * (a - 1) / 2) # baselines @@ -218,29 +219,18 @@ def test_simple_visibility_dataset(): ), ) - m = astropy_encode_msgpack(VISIBILITY) - ds_out = astropy_decode_msgpack(m) + message = astropy_encode_msgpack(VISIBILITY) + assert isinstance(message, bytes) + ds_out = astropy_decode_msgpack(message) xr.testing.assert_identical(VISIBILITY, ds_out) - # import zarr.storage as zs - - # store = zs.MemoryStore() - - # VISIBILITY.astropy.dequantify().to_zarr(store, mode="w", consolidated=True) - # dt_store = xr.open_datatree( - # store, engine="zarr", consolidated=True - # ).astropy.quantify() - - # xr.testing.assert_identical(VISIBILITY, dt_store) - - # expected = dataset_to_skycoord(VISIBILITY.field_phase_center.dataset) - # actual = dataset_to_skycoord(dt_store.field_phase_center.dataset) - # np.testing.assert_array_equal(actual, expected, strict=True) + assert dataset_to_skycoord(ds_out["field_phase_center"]) == ( + SkyCoord(ra=[0.1] * u.deg, dec=[0.5] * u.deg, frame=ICRS()) + ) -@requires_zarr @pytest.mark.parametrize( - "skycoords,frame,unit", + ("skycoords", "frame", "unit"), [ ( FIELD_PHASE_CENTER_ICRS, @@ -276,11 +266,10 @@ def test_visibility_dataset(skycoords: list[SkyCoord], frame, unit): 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( + 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") ), @@ -319,20 +308,6 @@ def test_visibility_dataset(skycoords: list[SkyCoord], frame, unit): children=dict( field_phase_center=xr.DataTree( skycoord_to_dataset( - # SkyCoord( - # SphericalRepresentation( - # lon=[[0.1, 0.2]] * u.deg, - # lat=[[0.5, 0.2]] * u.deg, - # distance=[[0, 0]] * u.pc, - # ).with_differentials( - # SphericalDifferential( - # d_lon=[[1, 1]] * u.deg / u.s, - # d_lat=[[1, 1]] * u.deg / u.s, - # d_distance=[[1, 1]] * u.pc / u.yr, - # ) - # ), - # frame=ICRS() - # ), SkyCoord( ra=[[0.1, 0.2]] * u.deg, dec=[[0.5, 0.2]] * u.deg, @@ -343,8 +318,8 @@ def test_visibility_dataset(skycoords: list[SkyCoord], frame, unit): frame=ICRS(), ), coords=[ - ("phase_center_label", [0]), - ("time_poly", [0, 1]), + ("phase_center_label", np.array([0])), + ("time_poly", np.array([0, 1])), ], ) ), @@ -357,7 +332,7 @@ def test_visibility_dataset(skycoords: list[SkyCoord], frame, unit): ), coords=[ ("calibrator_label", ["A", "B", "C", "D", "E"]), - ("time_poly", [0]), + ("time_poly", np.array([0])), ], ) ), @@ -367,21 +342,3 @@ def test_visibility_dataset(skycoords: list[SkyCoord], frame, unit): m = astropy_encode_msgpack(VISIBILITY) ds_out = astropy_decode_msgpack(m) xr.testing.assert_identical(VISIBILITY, ds_out) - - # import zarr.storage as zs - - # store = zs.MemoryStore() - - # save_datatree_compress_multi_index( - # VISIBILITY.astropy.dequantify(), store, mode="w", consolidated=True - # ) - # dt_store = open_datatree_decompress_multi_index( - # store, engine="zarr" - # ).astropy.quantify() - - # xr.testing.assert_identical(VISIBILITY, dt_store) - - # expected = dataset_to_skycoord(VISIBILITY.field_phase_center.dataset) - # actual = dataset_to_skycoord(dt_store.field_phase_center.dataset) - - # np.testing.assert_array_equal(actual, expected, strict=True) diff --git a/astropy_xarray/tests/utils.py b/astropy_xarray/tests/utils.py index 64d7e496..08e306e2 100644 --- a/astropy_xarray/tests/utils.py +++ b/astropy_xarray/tests/utils.py @@ -30,7 +30,6 @@ def importorskip(name): has_dask_array, requires_dask_array = importorskip("dask.array") has_scipy, requires_scipy = importorskip("scipy") has_bottleneck, requires_bottleneck = importorskip("bottleneck") -has_zarr, requires_zarr = importorskip("zarr") @contextmanager diff --git a/docs/examples/skycoord.ipynb b/docs/examples/skycoord.ipynb index 1b0c0411..120e37a4 100644 --- a/docs/examples/skycoord.ipynb +++ b/docs/examples/skycoord.ipynb @@ -4,7 +4,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## SkyCoord, Frame and Angle as Dataset\n", + "## SkyCoord, Frame and Angles as xarray Dataset\n", "\n", "SkyCoord and Frame fully support conversion to datasets preserving:\n", "* Representation\n", @@ -12,7 +12,7 @@ "* Frame\n", "* Angle (Longitude and Latitude)\n", "\n", - "Conversion to dataset supports optionally providing axis labels." + "Conversion to dataset supports optionally providing axis (xarray coordinate) labels." ] }, { @@ -27,7 +27,6 @@ "from astropy.coordinates import ICRS, SkyCoord\n", "\n", "from astropy_xarray.coordinates.sky_coord import (\n", - " DatasetRepresentation,\n", " skycoord_to_dataset,\n", ")\n", "\n", @@ -37,7 +36,7 @@ " 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: for correctness use:\n", + " # NOTE: to explicity be a direction, use:\n", " # frame=ICRS(\n", " # representation_type='unitspherical',\n", " # differential_type='unitsphericalcoslat',\n", @@ -51,7 +50,10 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "SkyCoordinate representation and differential types for display and data are also detected." + "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'`" ] }, { @@ -69,7 +71,21 @@ " frame=\"icrs\",\n", ")\n", "display(sky_position)\n", - "skycoord_to_dataset(sky_position, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])])" + "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\"])]))" ] }, { @@ -175,7 +191,7 @@ "source": [ "# SkyCoord itself doesn't strongly infer directional reprentation and differential types.\n", "import astropy.units as u\n", - "from astropy.coordinates import ICRS, SkyCoord\n", + "from astropy.coordinates import SkyCoord\n", "\n", "sky_direction = SkyCoord(\n", " ra=[2, 6, 7, 4] * u.deg,\n", @@ -258,9 +274,7 @@ "print(\"SkyCoord:\", sky_direction.ra.data[1])\n", "\n", "# dataset serialization in FRAME mode avoids unit conversions where possible\n", - "ds = skycoord_to_dataset(\n", - " sky_direction, DatasetRepresentation.FRAME, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])]\n", - ")\n", + "ds = skycoord_to_dataset(sky_direction, coords=[(\"field\", [\"a\", \"b\", \"c\", \"d\"])])\n", "print(\"Dataset:\", ds.ra.data[1].value)\n", "\n", "\n", diff --git a/pyproject.toml b/pyproject.toml index 83815a00..08c9a738 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,15 +18,16 @@ classifiers = [ "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python", + "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Topic :: Scientific/Engineering", ] -requires-python = ">=3.11" +requires-python = ">=3.10" dependencies = [ "numpy>=1.23,<2.3.0", "xarray>=2022.06.0,<=2025.4.0", - "astropy>=7.0.0,<7.2.0", + "astropy>=6.1.0,<7.2.0", "pandas>=2.3.0,<3.0.0", ] dynamic = ["version"] @@ -40,7 +41,6 @@ test = [ "pytest>=8.0", "pytest-cov", "dask[array]", - "zarr>=3.0.0; python_version >= '3.11'", "msgpack>=1.1.2", "msgpack_numpy>=0.4.8", "cf_xarray>=0.10.6", @@ -110,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"] From 1f13065550d0176cd8f03f06bfa1533221382c62 Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Fri, 27 Mar 2026 16:46:55 +0800 Subject: [PATCH 08/10] fix astropy 6.1 skycoord support Signed-off-by: Callan Gray --- astropy_xarray/coordinates/frame.py | 8 +++---- astropy_xarray/tests/test_sky_coord.py | 22 ++++++++----------- .../tests/test_visibility_serialization.py | 7 ------ 3 files changed, 13 insertions(+), 24 deletions(-) diff --git a/astropy_xarray/coordinates/frame.py b/astropy_xarray/coordinates/frame.py index 11341c54..54f3436e 100644 --- a/astropy_xarray/coordinates/frame.py +++ b/astropy_xarray/coordinates/frame.py @@ -57,13 +57,13 @@ def dump_frame(frame: BaseCoordinateFrame, with_data: bool = False) -> dict: ser = { "name": frame.name, - "representation_type": frame.representation_type.name, - "differential_type": frame.differential_type.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.name} + ser["data"] = {"representation_type": frame.data.get_name()} if "s" in frame.data.differentials: - ser["data"]["differential_type"] = frame.data.differentials["s"].name + 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) diff --git a/astropy_xarray/tests/test_sky_coord.py b/astropy_xarray/tests/test_sky_coord.py index 5713d0a3..b0c9d42f 100644 --- a/astropy_xarray/tests/test_sky_coord.py +++ b/astropy_xarray/tests/test_sky_coord.py @@ -1,5 +1,3 @@ -import sys - import astropy.units as u import numpy as np import pytest @@ -56,11 +54,6 @@ _skycoord_representation_component_names, ) -pytestmark = pytest.mark.skipif( - sys.version_info < (3, 11), - reason="python3.10 or higher required for skycoords support", -) - @pytest.mark.parametrize( ("expected_base", "expected_s"), @@ -302,15 +295,18 @@ def test_skycoord_roundtrip( ) ) ) - expected_frame_data_rep_name = frame.default_representation.name - expected_frame_data_diff_name = frame.default_differential.name + 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.name - expected_frame_data_diff_name = differential_type.name + expected_frame_data_rep_name = representation_type.get_name() + expected_frame_data_diff_name = differential_type.get_name() - assert expected.frame.data.name == expected_frame_data_rep_name - assert expected.frame.data.differentials["s"].name == expected_frame_data_diff_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"]), diff --git a/astropy_xarray/tests/test_visibility_serialization.py b/astropy_xarray/tests/test_visibility_serialization.py index 3eef7035..0ba96cb0 100644 --- a/astropy_xarray/tests/test_visibility_serialization.py +++ b/astropy_xarray/tests/test_visibility_serialization.py @@ -1,5 +1,4 @@ import json -import sys import astropy.units as u import msgpack @@ -20,12 +19,6 @@ skycoord_to_dataset, ) -pytestmark = pytest.mark.skipif( - sys.version_info < (3, 11), - reason="python3.10 or higher required for skycoords support", -) - - FIELD_PHASE_CENTER_ICRS = [ SkyCoord( ra=[2, 0.2, 0.02, 0.002] * u.deg, From 3be867644853eb3397d5be1ee3b31aa7ee4e79aa Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Fri, 27 Mar 2026 17:04:05 +0800 Subject: [PATCH 09/10] add api docs Signed-off-by: Callan Gray --- docs/api.rst | 5 +++++ 1 file changed, 5 insertions(+) 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:: From 6648a9a1dc2852a00459c02b5987c5c7e967f778 Mon Sep 17 00:00:00 2001 From: Callan Gray Date: Fri, 27 Mar 2026 17:28:36 +0800 Subject: [PATCH 10/10] remove use_frame_names, update changelog Signed-off-by: Callan Gray --- astropy_xarray/coordinates/sky_coord.py | 4 ++-- astropy_xarray/tests/test_frame_conversion.py | 8 -------- docs/examples/plotting.ipynb | 5 ----- docs/whats-new.rst | 6 ++++++ 4 files changed, 8 insertions(+), 15 deletions(-) diff --git a/astropy_xarray/coordinates/sky_coord.py b/astropy_xarray/coordinates/sky_coord.py index b8c8e30a..6a1c2472 100644 --- a/astropy_xarray/coordinates/sky_coord.py +++ b/astropy_xarray/coordinates/sky_coord.py @@ -100,7 +100,7 @@ def skycoord_to_dataset( ) -def dataset_to_skycoord(ds: xr.Dataset, use_frame_names: bool = True) -> SkyCoord: +def dataset_to_skycoord(ds: xr.Dataset) -> SkyCoord: """Convert a SkyCoord-based Dataset with metadata attributes to a SkyCoord. Args: @@ -115,7 +115,7 @@ def dataset_to_skycoord(ds: xr.Dataset, use_frame_names: bool = True) -> SkyCoor frame._data = load_representation( ds.attrs["frame"]["data"]["representation_type"], ds.attrs["frame"]["data"].get("differential_type"), - ds.attrs["frame"]["name"] if use_frame_names else None, + ds.attrs["frame"]["name"], {k: v.data for k, v in dsq.data_vars.items()}, ) frame.representation_type = ds.attrs["frame"]["representation_type"] diff --git a/astropy_xarray/tests/test_frame_conversion.py b/astropy_xarray/tests/test_frame_conversion.py index 10c13ff9..77945eaf 100644 --- a/astropy_xarray/tests/test_frame_conversion.py +++ b/astropy_xarray/tests/test_frame_conversion.py @@ -1,15 +1,11 @@ -from importlib.metadata import version - import astropy.units as u import numpy.testing -import pytest from astropy.coordinates import ( CartesianDifferential, CartesianRepresentation, SkyCoord, ) from astropy.coordinates.builtin_frames import ICRS -from packaging.version import Version from astropy_xarray.coordinates import ( skycoord_to_dataset, @@ -44,10 +40,6 @@ def test_frame_conversion_aliasing(): numpy.testing.assert_array_equal(sc.pm_ra_cosdec, sc2.pm_ra_cosdec) -@pytest.mark.xfail( - condition=Version(version("astropy")) <= Version("7.0.0"), - reason="astropy 6.0.0 skycoords not supported", -) def test_frame_conversion_components(): sc = SkyCoord( ICRS().realize_frame( diff --git a/docs/examples/plotting.ipynb b/docs/examples/plotting.ipynb index b644025c..25d95db6 100644 --- a/docs/examples/plotting.ipynb +++ b/docs/examples/plotting.ipynb @@ -240,10 +240,6 @@ ")\n", "\n", "\n", - "s1 = skycoord_to_dataset(sc, coords=[(\"field\", [0, 1, 2, 3])], use_frame_names=False)\n", - "s2 = skycoord_to_dataset(\n", - " sc, coords=[(\"field\", [0, 1, 2, 3])], use_frame_names=True, skip_implicit=False\n", - ")\n", "s = skycoord_to_dataset(sc, coords=[(\"field\", [0, 1, 2, 3])])\n", "\n", "# display(s1)\n", @@ -258,7 +254,6 @@ "metadata": {}, "outputs": [], "source": [ - "# display(dataset_to_skycoord(s1, use_frame_names=False))\n", "display(dataset_to_skycoord(s))" ] }, 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) -----------------