diff --git a/docs/blocks/transform.md b/docs/blocks/transform.md index 066706b..be131a0 100644 --- a/docs/blocks/transform.md +++ b/docs/blocks/transform.md @@ -1 +1,178 @@ +--- +jupytext: + text_representation: + format_name: myst +kernelspec: + display_name: Python 3 + name: python +--- + # Functional transformations with `CoordinateTransformIndex` + +## Highlights + +1. The coordinate variables whose values are described by _coordinate + transformations_ -- i.e., by a set of formulas describing the relationship + between array indices and coordinate labels -- are best handled via an + {py:class}`xarray.indexes.CoordinateTransformIndex`. +1. Coordinate variables associated with such index are lazy and therefore + use very little memory. They may have arbitrary dimensions. +1. Alignment may be implemented in an optimal way, i.e., based on coordinate transformation + parameters rather than on raw coordinate labels. +1. Xarray exposes an abstract {py:class}`~xarray.indexes.CoordinateTransform` + class to plug in 3rd-party coordinate transformations with support + of dimension and coordinate variable names (see the example below). + +```{seealso} +`CoordinateTransformIndex` is often used as a building block by other +custom indexes such as {py:class}`xarray.indexes.RangeIndex` (see +{doc}`../builtin/range`) and {py:class}`rasterix.RasterIndex` (see +{doc}`../earth/raster`). +``` + +## Example (Astronomy) + +As a real-world example, let's create a custom +{py:class}`xarray.indexes.CoordinateTransform` that wraps an +{py:class}`astropy.wcs.WCS` object. This Xarray coordinate transform adapter +class simply maps Xarray dimension and coordinate variable names to pixel +and world axis names of the [shared Python interface for World Coordinate +Systems](https://doi.org/10.5281/zenodo.1188874) used in Astropy. + +```{note} +This example is taken and adapted from +[this gist](https://gist.github.com/Cadair/4a03750868e044ac4bdd6f3a04ed7abc) by +[Stuart Mumford](https://github.com/Cadair). + +It only provides basic integration between Astropy's WCS and Xarray +coordinate transforms. More advanced integration could leverage the +[slicing capabilities of WCS objects](https://docs.astropy.org/en/latest/wcs/wcsapi.html#slicing-of-wcs-objects) +in a custom {py:class}`xarray.indexes.CoordinateTransformIndex` subclass. + +``` + +```{code-cell} python +from collections.abc import Hashable +from typing import Any + +import numpy as np +import xarray as xr +from astropy.wcs import WCS + + +def escape(name): + return name.replace(".", "_").replace("custom:", "") + + +class WCSCoordinateTransform(xr.indexes.CoordinateTransform): + """Lightweight adapter class for the World Coordinate Systems (WCS) API. + + More info: https://docs.astropy.org/en/latest/wcs/wcsapi.html + """ + def __init__(self, wcs: WCS): + pixel_axis_names = [ + pan or f"dim{i}" + for i, pan in enumerate(wcs.pixel_axis_names) + ] + world_axis_names = [ + escape(wan or wphy or f"coord{i}") + for i, (wan, wphy) in enumerate( + zip(wcs.world_axis_names, wcs.world_axis_physical_types) + ) + ] + dim_size = { + name: size + for name, size in zip(pixel_axis_names, wcs.array_shape) + } + + super().__init__(world_axis_names, dim_size, dtype=np.dtype(float)) + self.wcs = wcs + + def forward(self, dim_positions: dict[str, Any]) -> dict[Hashable, Any]: + """Perform array -> world coordinate transformation.""" + + pixel = [dim_positions[dim] for dim in self.dims] + world = self.wcs.array_index_to_world_values(*pixel) + + return {name: w for name, w in zip(self.coord_names, world)} + + def reverse(self, coord_labels: dict[Hashable, Any]) -> dict[str, Any]: + """Perform world -> array coordinate reverse transformation.""" + + world = [coord_labels[name] for name in self.coord_names] + pixel = self.wcs.world_to_array_index_values(*world) + + return {name: p for name, p in zip(self.dims, pixel)} +``` + +```{code-cell} python +--- +tags: [remove-cell] +--- +%xmode minimal +xr.set_options( + display_expand_indexes=True, + display_expand_data=True, +); +np.set_printoptions(precision=3, threshold=10, edgeitems=2); +``` + +### Assigning + +Let's now create a small function that opens a FITS file with Astropy, creates +an Xarray {py:class}`~xarray.indexes.CoordinateTransformIndex` and its +associated lazy coordinate variables from the {py:class}`~astropy.wcs.WCS` +object and returns both the data and coordinates as an +{py:class}`xarray.DataArray`. + +```{code-cell} python +from astropy.io import fits +from astropy.utils.data import get_pkg_data_filename + + +def open_fits_dataarray(filename, item=0): + hdu = fits.open(filename)[item] + wcs = WCS(hdu.header) + + transform = WCSCoordinateTransform(wcs) + index = xr.indexes.CoordinateTransformIndex(transform) + coords = xr.Coordinates.from_xindex(index) + + return xr.DataArray( + hdu.data, coords=coords, dims=transform.dims, attrs={"wcs": wcs} + ) + +``` + +Open a simple image with two celestial axes. + +```{code-cell} python +fname = get_pkg_data_filename("galactic_center/gc_2mass_k.fits") + +da_2d = open_fits_dataarray(fname) +da_2d +``` + +```{code-cell} python +# lazy coordinate variables! + +da_2d.pos_eq_ra +``` + +```{code-cell} python +da_2d.plot.pcolormesh( + x="pos_eq_ra", + y="pos_eq_dec", + vmax=1300, + cmap="magma", +); +``` + +Open a spectral cube with two celestial axes and one spectral axis. + +```{code-cell} python +fname = get_pkg_data_filename("l1448/l1448_13co.fits") + +da_3d = open_fits_dataarray(fname) +da_3d +``` diff --git a/docs/builtin/pdinterval.md b/docs/builtin/pdinterval.md index e55025d..bf6949f 100644 --- a/docs/builtin/pdinterval.md +++ b/docs/builtin/pdinterval.md @@ -26,7 +26,7 @@ Learn more at the [Pandas](https://pandas.pydata.org/pandas-docs/stable/user_gui ``` ```` -# Highlights +## Highlights 1. Xarray's built-in support for pandas Index classes extends to more sophisticated classes like {py:class}`pandas.IntervalIndex`. 1. Xarray now generates such indexes automatically when using {py:meth}`xarray.DataArray.groupby_bins` or {py:meth}`xarray.Dataset.groupby_bins`. diff --git a/docs/builtin/range.md b/docs/builtin/range.md index a0a46e8..0a55995 100644 --- a/docs/builtin/range.md +++ b/docs/builtin/range.md @@ -15,8 +15,9 @@ kernelspec: ranges. Fortunately, there is {py:class}`xarray.indexes.RangeIndex` that works with real numbers. 1. Xarray's `RangeIndex` is built on top of - {py:class}`xarray.indexes.CoordinateTransformIndex` and therefore supports - very large ranges represented as lazy coordinate variables. + {py:class}`xarray.indexes.CoordinateTransformIndex` (see + {doc}`../blocks/transform`) and therefore supports very large ranges + represented as lazy coordinate variables. ## Example diff --git a/docs/conf.py b/docs/conf.py index 0c428b3..bc3fa40 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -124,4 +124,5 @@ "geopandas": ("https://geopandas.org/en/stable/", None), "pint-xarray": ("https://pint-xarray.readthedocs.io/en/latest/", None), "pint": ("https://pint.readthedocs.io/en/stable/", None), + "astropy": ("https://docs.astropy.org/en/latest/", None), } diff --git a/docs/earth/raster.md b/docs/earth/raster.md index bd9f594..bef296d 100644 --- a/docs/earth/raster.md +++ b/docs/earth/raster.md @@ -31,7 +31,7 @@ Learn more at the [Rasterix](https://rasterix.readthedocs.io/en/latest/) documen Rasterix provides a RasterIndex that allows indexing using a functional transformation defined by an _affine transform_. -It uses {py:class}`~xarray.indexes.CoordinateTransformIndex` as a building block. In doing so, +It uses {py:class}`~xarray.indexes.CoordinateTransformIndex` as a building block (see {doc}`../blocks/transform`). In doing so, 1. RasterIndex eliminates an entire class of bugs where Xarray allows you to add (for example) two datasets with different affine transforms (and/or projections) and return nonsensical outputs. 1. The associated coordinate variables are lazy, and use very little memory. Thus very large coordinate frames can be represented. diff --git a/requirements.txt b/requirements.txt index 5688e1a..3dea191 100644 --- a/requirements.txt +++ b/requirements.txt @@ -21,4 +21,5 @@ xvec git+https://github.com/dcherian/rolodex pint-xarray cf_xarray +astropy git+https://github.com/pydata/xarray