|
| 1 | +--- |
| 2 | +jupytext: |
| 3 | + text_representation: |
| 4 | + format_name: myst |
| 5 | +kernelspec: |
| 6 | + display_name: Python 3 |
| 7 | + name: python |
| 8 | +--- |
| 9 | + |
1 | 10 | # 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 | +``` |
0 commit comments