Add coordinate transform index (#24)

* nit

* add astropy WCS adapter

* clean-up

* add cross-refs

* add more description

* improve figure constrast and color

* show lazy coordinate variable

* example title

* tweaks

---------

Co-authored-by: Deepak Cherian <deepak@cherian.net>
Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

Author: benbovy

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
+```

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`.

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
 

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),
 }

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.

requirements.txt

@@ -21,4 +21,5 @@ xvec
 git+https://github.com/dcherian/rolodex
 pint-xarray
 cf_xarray
+astropy
 git+https://github.com/pydata/xarray