-
Notifications
You must be signed in to change notification settings - Fork 3
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
benbovy
merged 11 commits into
xarray-contrib:main
from
benbovy:add-coordinate-transform-index
Jul 9, 2025
Merged
Changes from all commits
Commits
Show all changes
11 commits
Select commit
Hold shift + click to select a range
e935003
nit
benbovy b386c4a
add astropy WCS adapter
benbovy d92461d
clean-up
benbovy 019a6e2
Merge branch 'main' into add-coordinate-transform-index
benbovy 4967c28
add cross-refs
benbovy 529cd7b
add more description
benbovy 6659547
improve figure constrast and color
benbovy 564859d
show lazy coordinate variable
benbovy 263ff56
example title
benbovy b10c1a9
tweaks
dcherian 4ce5a0c
Merge branch 'main' into add-coordinate-transform-index
dcherian File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
``` |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -21,4 +21,5 @@ xvec | |
git+https://github.com/dcherian/rolodex | ||
pint-xarray | ||
cf_xarray | ||
astropy | ||
git+https://github.com/pydata/xarray |
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.