Skip to content

Add coordinate transform index #24

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Jul 9, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
177 changes: 177 additions & 0 deletions docs/blocks/transform.md
Original file line number Diff line number Diff line change
@@ -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
```
2 changes: 1 addition & 1 deletion docs/builtin/pdinterval.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down
5 changes: 3 additions & 2 deletions docs/builtin/range.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
1 change: 1 addition & 0 deletions docs/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}
2 changes: 1 addition & 1 deletion docs/earth/raster.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -21,4 +21,5 @@ xvec
git+https://github.com/dcherian/rolodex
pint-xarray
cf_xarray
astropy
git+https://github.com/pydata/xarray