From e935003888715c753162a5a373d487ed115589a2 Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Wed, 9 Jul 2025 02:11:27 +0200 Subject: [PATCH 1/9] nit --- docs/builtin/pdinterval.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/builtin/pdinterval.md b/docs/builtin/pdinterval.md index 2d1875b..a2df391 100644 --- a/docs/builtin/pdinterval.md +++ b/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`. From b386c4a521ea8d2ff853ac64bcbdacdcaec38a56 Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Wed, 9 Jul 2025 02:11:41 +0200 Subject: [PATCH 2/9] add astropy WCS adapter --- docs/blocks/transform.md | 162 +++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + 2 files changed, 163 insertions(+) diff --git a/docs/blocks/transform.md b/docs/blocks/transform.md index 066706b..a7599e6 100644 --- a/docs/blocks/transform.md +++ b/docs/blocks/transform.md @@ -1 +1,163 @@ +--- +jupytext: + text_representation: + format_name: myst +kernelspec: + display_name: Python 3 + name: python +--- + # Functional transformations with `CoordinateTransformIndex` + +## Highlights + +## Example + +Taken and adapted from +[here](https://gist.github.com/Cadair/4a03750868e044ac4bdd6f3a04ed7abc) (author: +Stuart Mumford). + +```{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_attrs=False, +); +``` + +```{code-cell} python +from astropy.io import fits +from astropy.utils.data import get_pkg_data_filename + + +filename = get_pkg_data_filename('l1448/l1448_13co.fits') +hdu = fits.open(filename)[0] +wcs = WCS(hdu.header) +data = hdu.data + +transform = WCSCoordinateTransform(wcs) +index = xr.indexes.CoordinateTransformIndex(transform) +coords = xr.Coordinates.from_xindex(index) + +ds = xr.Dataset({'primary': (transform.dims, data)}, coords=coords) +ds +``` + +```{code-cell} python +da_subset = ds.primary.isel(dim2=[45, 50, 55]) + +da_subset.plot.pcolormesh( + x="pos_eq_ra", + y="spect_dopplerVeloc_opt", + col="dim2", + col_wrap=2, +); +``` + +```{code-cell} python +sliced = ds.isel(dim2=50) + +ds.primary.isel(dim2=50).plot(subplot_kws=dict(projection=wcs, slices=(50, 'x', 'y'))) +``` + +```{code-cell} python +filename2 = get_pkg_data_filename('galactic_center/gc_2mass_k.fits') +hdu2 = fits.open(filename2)[0] +wcs2 = WCS(hdu2.header) +data2 = hdu2.data + +transform2 = WCSCoordinateTransform(wcs2) +index2 = xr.indexes.CoordinateTransformIndex(transform2) +coords2 = xr.Coordinates.from_xindex(index2) + +ds2 = xr.Dataset({'primary': (transform2.dims, data2)}, coords=coords2) +ds2 +``` + +```{code-cell} python +ds2.primary.plot.pcolormesh(x="pos_eq_ra", y="pos_eq_dec"); +``` + +```{code-cell} python +hdu2 = fits.open("http://data.sunpy.org/sunpy/v1/AIA20110607_063302_0171_lowres.fits")[1] +wcs2 = WCS(hdu2.header) +data2 = hdu2.data + +transform2 = WCSCoordinateTransform(wcs2) +index2 = xr.indexes.CoordinateTransformIndex(transform2) +coords2 = xr.Coordinates.from_xindex(index2) + +ds2 = xr.Dataset({'data': (transform2.dims, data2)}, coords=coords2) +ds2 +``` + +```{code-cell} python +ds2.data.plot.pcolormesh(x="pos_helioprojective_lon", y="pos_helioprojective_lat"); +``` + +```{code-cell} python +ds2.data.plot(subplot_kws=dict(projection=wcs2), vmin=0, vmax=2000) +``` diff --git a/requirements.txt b/requirements.txt index cd3b716..b3df226 100644 --- a/requirements.txt +++ b/requirements.txt @@ -21,3 +21,4 @@ xvec git+https://github.com/dcherian/rolodex pint-xarray cf_xarray +astropy From d92461d0c24ebd2c80e6c00a5f044315f180da94 Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Wed, 9 Jul 2025 09:34:17 +0200 Subject: [PATCH 3/9] clean-up --- docs/blocks/transform.md | 69 ++++++++++------------------------------ 1 file changed, 16 insertions(+), 53 deletions(-) diff --git a/docs/blocks/transform.md b/docs/blocks/transform.md index a7599e6..e4b0854 100644 --- a/docs/blocks/transform.md +++ b/docs/blocks/transform.md @@ -93,71 +93,34 @@ from astropy.io import fits from astropy.utils.data import get_pkg_data_filename -filename = get_pkg_data_filename('l1448/l1448_13co.fits') -hdu = fits.open(filename)[0] -wcs = WCS(hdu.header) -data = hdu.data +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) + transform = WCSCoordinateTransform(wcs) + index = xr.indexes.CoordinateTransformIndex(transform) + coords = xr.Coordinates.from_xindex(index) -ds = xr.Dataset({'primary': (transform.dims, data)}, coords=coords) -ds -``` - -```{code-cell} python -da_subset = ds.primary.isel(dim2=[45, 50, 55]) - -da_subset.plot.pcolormesh( - x="pos_eq_ra", - y="spect_dopplerVeloc_opt", - col="dim2", - col_wrap=2, -); -``` - -```{code-cell} python -sliced = ds.isel(dim2=50) + return xr.DataArray( + hdu.data, coords=coords, dims=transform.dims, attrs={"wcs": wcs} + ) -ds.primary.isel(dim2=50).plot(subplot_kws=dict(projection=wcs, slices=(50, 'x', 'y'))) ``` ```{code-cell} python -filename2 = get_pkg_data_filename('galactic_center/gc_2mass_k.fits') -hdu2 = fits.open(filename2)[0] -wcs2 = WCS(hdu2.header) -data2 = hdu2.data +fname = get_pkg_data_filename("galactic_center/gc_2mass_k.fits") -transform2 = WCSCoordinateTransform(wcs2) -index2 = xr.indexes.CoordinateTransformIndex(transform2) -coords2 = xr.Coordinates.from_xindex(index2) - -ds2 = xr.Dataset({'primary': (transform2.dims, data2)}, coords=coords2) -ds2 -``` - -```{code-cell} python -ds2.primary.plot.pcolormesh(x="pos_eq_ra", y="pos_eq_dec"); +da_2d = open_fits_dataarray(fname) +da_2d ``` ```{code-cell} python -hdu2 = fits.open("http://data.sunpy.org/sunpy/v1/AIA20110607_063302_0171_lowres.fits")[1] -wcs2 = WCS(hdu2.header) -data2 = hdu2.data - -transform2 = WCSCoordinateTransform(wcs2) -index2 = xr.indexes.CoordinateTransformIndex(transform2) -coords2 = xr.Coordinates.from_xindex(index2) - -ds2 = xr.Dataset({'data': (transform2.dims, data2)}, coords=coords2) -ds2 +da_2d.plot.pcolormesh(x="pos_eq_ra", y="pos_eq_dec"); ``` ```{code-cell} python -ds2.data.plot.pcolormesh(x="pos_helioprojective_lon", y="pos_helioprojective_lat"); -``` +fname = get_pkg_data_filename("l1448/l1448_13co.fits") -```{code-cell} python -ds2.data.plot(subplot_kws=dict(projection=wcs2), vmin=0, vmax=2000) +da_3d = open_fits_dataarray(fname) +da_3d ``` From 4967c28b50909a9741e1710d98ad70b6d4c21a41 Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Wed, 9 Jul 2025 11:20:53 +0200 Subject: [PATCH 4/9] add cross-refs --- docs/builtin/range.md | 5 +++-- docs/conf.py | 1 + docs/earth/raster.md | 2 +- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/builtin/range.md b/docs/builtin/range.md index a0a46e8..0a55995 100644 --- a/docs/builtin/range.md +++ b/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 diff --git a/docs/conf.py b/docs/conf.py index b109771..f270542 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -124,4 +124,5 @@ "geopandas": ("https://geopandas.readthedocs.io/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), } diff --git a/docs/earth/raster.md b/docs/earth/raster.md index bd9f594..bef296d 100644 --- a/docs/earth/raster.md +++ b/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. From 529cd7b001da7a6389b76c6c5cbe613fdbd6855c Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Wed, 9 Jul 2025 11:21:08 +0200 Subject: [PATCH 5/9] add more description --- docs/blocks/transform.md | 50 +++++++++++++++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 3 deletions(-) diff --git a/docs/blocks/transform.md b/docs/blocks/transform.md index e4b0854..c8f2311 100644 --- a/docs/blocks/transform.md +++ b/docs/blocks/transform.md @@ -11,11 +11,43 @@ kernelspec: ## 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. The coordinate variables associated with such index are lazy and therefore + use very little memory. They may have arbitrary dimensions. +1. The alignment of datasets and dataarrays 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). +1. `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 -Taken and adapted from -[here](https://gist.github.com/Cadair/4a03750868e044ac4bdd6f3a04ed7abc) (author: -Stuart Mumford). +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. + +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 @@ -88,6 +120,14 @@ xr.set_options( ); ``` +### 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 @@ -107,6 +147,8 @@ def open_fits_dataarray(filename, item=0): ``` +Open a simple image with two celestial axes. + ```{code-cell} python fname = get_pkg_data_filename("galactic_center/gc_2mass_k.fits") @@ -118,6 +160,8 @@ da_2d da_2d.plot.pcolormesh(x="pos_eq_ra", y="pos_eq_dec"); ``` +Open a spectral cube with two celestial axes and one spectral axis. + ```{code-cell} python fname = get_pkg_data_filename("l1448/l1448_13co.fits") From 6659547e7c172b22b57f5d238563e262f3d4f43b Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Wed, 9 Jul 2025 12:05:14 +0200 Subject: [PATCH 6/9] improve figure constrast and color --- docs/blocks/transform.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/blocks/transform.md b/docs/blocks/transform.md index c8f2311..4f47d5e 100644 --- a/docs/blocks/transform.md +++ b/docs/blocks/transform.md @@ -157,7 +157,12 @@ da_2d ``` ```{code-cell} python -da_2d.plot.pcolormesh(x="pos_eq_ra", y="pos_eq_dec"); +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. From 564859dfc18b5881f54283643147daf108e7a6dc Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Wed, 9 Jul 2025 12:15:39 +0200 Subject: [PATCH 7/9] show lazy coordinate variable --- docs/blocks/transform.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/docs/blocks/transform.md b/docs/blocks/transform.md index 4f47d5e..bae98d9 100644 --- a/docs/blocks/transform.md +++ b/docs/blocks/transform.md @@ -116,7 +116,7 @@ tags: [remove-cell] %xmode minimal xr.set_options( display_expand_indexes=True, - display_expand_attrs=False, + display_expand_data=True, ); ``` @@ -156,6 +156,12 @@ 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", From 263ff56f301d61b18a6ed4c0c4719f95f63006a3 Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Wed, 9 Jul 2025 12:32:29 +0200 Subject: [PATCH 8/9] example title --- docs/blocks/transform.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/blocks/transform.md b/docs/blocks/transform.md index bae98d9..0d3a171 100644 --- a/docs/blocks/transform.md +++ b/docs/blocks/transform.md @@ -28,7 +28,7 @@ kernelspec: {doc}`../builtin/range`) and {py:class}`rasterix.RasterIndex` (see {doc}`../earth/raster`). -## Example +## Example (Astronomy) As a real-world example, let's create a custom {py:class}`xarray.indexes.CoordinateTransform` that wraps an From b10c1a9f16dd50a2e0073f4ea43960a18c768be8 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Wed, 9 Jul 2025 07:20:33 -0700 Subject: [PATCH 9/9] tweaks --- docs/blocks/transform.md | 29 +++++++++++++---------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/docs/blocks/transform.md b/docs/blocks/transform.md index 0d3a171..be131a0 100644 --- a/docs/blocks/transform.md +++ b/docs/blocks/transform.md @@ -15,18 +15,20 @@ kernelspec: 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. The coordinate variables associated with such index are lazy and therefore +1. Coordinate variables associated with such index are lazy and therefore use very little memory. They may have arbitrary dimensions. -1. The alignment of datasets and dataarrays may be implemented in an optimal - way, i.e., based on coordinate transformation parameters rather than on - raw coordinate labels. +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). -1. `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`). + +```{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) @@ -40,7 +42,7 @@ 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. +[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 @@ -66,12 +68,8 @@ 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, - ): + def __init__(self, wcs: WCS): pixel_axis_names = [ pan or f"dim{i}" for i, pan in enumerate(wcs.pixel_axis_names) @@ -105,8 +103,6 @@ class WCSCoordinateTransform(xr.indexes.CoordinateTransform): pixel = self.wcs.world_to_array_index_values(*world) return {name: p for name, p in zip(self.dims, pixel)} - - ``` ```{code-cell} python @@ -118,6 +114,7 @@ xr.set_options( display_expand_indexes=True, display_expand_data=True, ); +np.set_printoptions(precision=3, threshold=10, edgeitems=2); ``` ### Assigning