Skip to content

Commit ac94b97

Browse files
benbovydcherian
andauthored
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>
1 parent 74110e2 commit ac94b97

File tree

6 files changed

+184
-4
lines changed

6 files changed

+184
-4
lines changed

docs/blocks/transform.md

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,178 @@
1+
---
2+
jupytext:
3+
text_representation:
4+
format_name: myst
5+
kernelspec:
6+
display_name: Python 3
7+
name: python
8+
---
9+
110
# Functional transformations with `CoordinateTransformIndex`
11+
12+
## Highlights
13+
14+
1. The coordinate variables whose values are described by _coordinate
15+
transformations_ -- i.e., by a set of formulas describing the relationship
16+
between array indices and coordinate labels -- are best handled via an
17+
{py:class}`xarray.indexes.CoordinateTransformIndex`.
18+
1. Coordinate variables associated with such index are lazy and therefore
19+
use very little memory. They may have arbitrary dimensions.
20+
1. Alignment may be implemented in an optimal way, i.e., based on coordinate transformation
21+
parameters rather than on raw coordinate labels.
22+
1. Xarray exposes an abstract {py:class}`~xarray.indexes.CoordinateTransform`
23+
class to plug in 3rd-party coordinate transformations with support
24+
of dimension and coordinate variable names (see the example below).
25+
26+
```{seealso}
27+
`CoordinateTransformIndex` is often used as a building block by other
28+
custom indexes such as {py:class}`xarray.indexes.RangeIndex` (see
29+
{doc}`../builtin/range`) and {py:class}`rasterix.RasterIndex` (see
30+
{doc}`../earth/raster`).
31+
```
32+
33+
## Example (Astronomy)
34+
35+
As a real-world example, let's create a custom
36+
{py:class}`xarray.indexes.CoordinateTransform` that wraps an
37+
{py:class}`astropy.wcs.WCS` object. This Xarray coordinate transform adapter
38+
class simply maps Xarray dimension and coordinate variable names to pixel
39+
and world axis names of the [shared Python interface for World Coordinate
40+
Systems](https://doi.org/10.5281/zenodo.1188874) used in Astropy.
41+
42+
```{note}
43+
This example is taken and adapted from
44+
[this gist](https://gist.github.com/Cadair/4a03750868e044ac4bdd6f3a04ed7abc) by
45+
[Stuart Mumford](https://github.com/Cadair).
46+
47+
It only provides basic integration between Astropy's WCS and Xarray
48+
coordinate transforms. More advanced integration could leverage the
49+
[slicing capabilities of WCS objects](https://docs.astropy.org/en/latest/wcs/wcsapi.html#slicing-of-wcs-objects)
50+
in a custom {py:class}`xarray.indexes.CoordinateTransformIndex` subclass.
51+
52+
```
53+
54+
```{code-cell} python
55+
from collections.abc import Hashable
56+
from typing import Any
57+
58+
import numpy as np
59+
import xarray as xr
60+
from astropy.wcs import WCS
61+
62+
63+
def escape(name):
64+
return name.replace(".", "_").replace("custom:", "")
65+
66+
67+
class WCSCoordinateTransform(xr.indexes.CoordinateTransform):
68+
"""Lightweight adapter class for the World Coordinate Systems (WCS) API.
69+
70+
More info: https://docs.astropy.org/en/latest/wcs/wcsapi.html
71+
"""
72+
def __init__(self, wcs: WCS):
73+
pixel_axis_names = [
74+
pan or f"dim{i}"
75+
for i, pan in enumerate(wcs.pixel_axis_names)
76+
]
77+
world_axis_names = [
78+
escape(wan or wphy or f"coord{i}")
79+
for i, (wan, wphy) in enumerate(
80+
zip(wcs.world_axis_names, wcs.world_axis_physical_types)
81+
)
82+
]
83+
dim_size = {
84+
name: size
85+
for name, size in zip(pixel_axis_names, wcs.array_shape)
86+
}
87+
88+
super().__init__(world_axis_names, dim_size, dtype=np.dtype(float))
89+
self.wcs = wcs
90+
91+
def forward(self, dim_positions: dict[str, Any]) -> dict[Hashable, Any]:
92+
"""Perform array -> world coordinate transformation."""
93+
94+
pixel = [dim_positions[dim] for dim in self.dims]
95+
world = self.wcs.array_index_to_world_values(*pixel)
96+
97+
return {name: w for name, w in zip(self.coord_names, world)}
98+
99+
def reverse(self, coord_labels: dict[Hashable, Any]) -> dict[str, Any]:
100+
"""Perform world -> array coordinate reverse transformation."""
101+
102+
world = [coord_labels[name] for name in self.coord_names]
103+
pixel = self.wcs.world_to_array_index_values(*world)
104+
105+
return {name: p for name, p in zip(self.dims, pixel)}
106+
```
107+
108+
```{code-cell} python
109+
---
110+
tags: [remove-cell]
111+
---
112+
%xmode minimal
113+
xr.set_options(
114+
display_expand_indexes=True,
115+
display_expand_data=True,
116+
);
117+
np.set_printoptions(precision=3, threshold=10, edgeitems=2);
118+
```
119+
120+
### Assigning
121+
122+
Let's now create a small function that opens a FITS file with Astropy, creates
123+
an Xarray {py:class}`~xarray.indexes.CoordinateTransformIndex` and its
124+
associated lazy coordinate variables from the {py:class}`~astropy.wcs.WCS`
125+
object and returns both the data and coordinates as an
126+
{py:class}`xarray.DataArray`.
127+
128+
```{code-cell} python
129+
from astropy.io import fits
130+
from astropy.utils.data import get_pkg_data_filename
131+
132+
133+
def open_fits_dataarray(filename, item=0):
134+
hdu = fits.open(filename)[item]
135+
wcs = WCS(hdu.header)
136+
137+
transform = WCSCoordinateTransform(wcs)
138+
index = xr.indexes.CoordinateTransformIndex(transform)
139+
coords = xr.Coordinates.from_xindex(index)
140+
141+
return xr.DataArray(
142+
hdu.data, coords=coords, dims=transform.dims, attrs={"wcs": wcs}
143+
)
144+
145+
```
146+
147+
Open a simple image with two celestial axes.
148+
149+
```{code-cell} python
150+
fname = get_pkg_data_filename("galactic_center/gc_2mass_k.fits")
151+
152+
da_2d = open_fits_dataarray(fname)
153+
da_2d
154+
```
155+
156+
```{code-cell} python
157+
# lazy coordinate variables!
158+
159+
da_2d.pos_eq_ra
160+
```
161+
162+
```{code-cell} python
163+
da_2d.plot.pcolormesh(
164+
x="pos_eq_ra",
165+
y="pos_eq_dec",
166+
vmax=1300,
167+
cmap="magma",
168+
);
169+
```
170+
171+
Open a spectral cube with two celestial axes and one spectral axis.
172+
173+
```{code-cell} python
174+
fname = get_pkg_data_filename("l1448/l1448_13co.fits")
175+
176+
da_3d = open_fits_dataarray(fname)
177+
da_3d
178+
```

docs/builtin/pdinterval.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ Learn more at the [Pandas](https://pandas.pydata.org/pandas-docs/stable/user_gui
2626
```
2727
````
2828

29-
# Highlights
29+
## Highlights
3030

3131
1. Xarray's built-in support for pandas Index classes extends to more sophisticated classes like {py:class}`pandas.IntervalIndex`.
3232
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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,9 @@ kernelspec:
1515
ranges. Fortunately, there is {py:class}`xarray.indexes.RangeIndex` that
1616
works with real numbers.
1717
1. Xarray's `RangeIndex` is built on top of
18-
{py:class}`xarray.indexes.CoordinateTransformIndex` and therefore supports
19-
very large ranges represented as lazy coordinate variables.
18+
{py:class}`xarray.indexes.CoordinateTransformIndex` (see
19+
{doc}`../blocks/transform`) and therefore supports very large ranges
20+
represented as lazy coordinate variables.
2021

2122
## Example
2223

docs/conf.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,4 +124,5 @@
124124
"geopandas": ("https://geopandas.org/en/stable/", None),
125125
"pint-xarray": ("https://pint-xarray.readthedocs.io/en/latest/", None),
126126
"pint": ("https://pint.readthedocs.io/en/stable/", None),
127+
"astropy": ("https://docs.astropy.org/en/latest/", None),
127128
}

docs/earth/raster.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ Learn more at the [Rasterix](https://rasterix.readthedocs.io/en/latest/) documen
3131

3232
Rasterix provides a RasterIndex that allows indexing using a functional transformation defined by an _affine transform_.
3333

34-
It uses {py:class}`~xarray.indexes.CoordinateTransformIndex` as a building block. In doing so,
34+
It uses {py:class}`~xarray.indexes.CoordinateTransformIndex` as a building block (see {doc}`../blocks/transform`). In doing so,
3535

3636
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.
3737
1. The associated coordinate variables are lazy, and use very little memory. Thus very large coordinate frames can be represented.

requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,4 +21,5 @@ xvec
2121
git+https://github.com/dcherian/rolodex
2222
pint-xarray
2323
cf_xarray
24+
astropy
2425
git+https://github.com/pydata/xarray

0 commit comments

Comments
 (0)