Idea: functionally-derived non-dimensional coordinates

@Cadair and I are from the solar and astrophysics communities, respectively (particularly SunPy and Astropy). In our fields, we have a concept of something called “World Coordinate Systems” (WCS) which basically are arbitrary mappings from pixel coordinates (which is often but not necessarily the same as the index) to physical coordinates. (For more on this and associated Python/Astropy APIs, see this document). If we are reading correctly, this concept maps roughly onto the xarray concept of “Non-dimension coordinates”.

However, a critical difference is this: WCS are usually expressed as arbitrary complex mathematical functions, rather than coordinate arrays, as it is crucial to some of the science cases to carry sub-pixel or other coordinate-related metadata around along with the WCS.

So our question is: is it in-scope for xarray non-dimensional coordinates to be extended to be functional instead of having to be arrays? I.e., have the coordinate arrays be generated on-the-fly from some function instead of being realized as arrays at creation-time. We have thought about several ways this might be specified and are willing to do some trial implementations, but are first asking here if it is likely to be

1. Easy
2. Hard
3. Impossible
4. PR will immediately be rejected on philosophical grounds, regardless?

Thanks!

[benbovy]

This looks like a nice use case for the forthcoming Xarray’s custom index feature.

How I see CRS/WCS-aware Xarray datasets with custom indexes:

• A set of coordinate(s) and their attributes hold data or metadata relevant for public use and that could be easily (de)serialized

• A custom index (CRSIndex or WCSIndex) provides CRS/WCS-aware implementations of common Xarray operations such as alignment (merge/concat) and data selection (sel), via Xarray.Index’s equals, union, intersection and query methods added in #5102 and #5322 (not yet ready for use outside of Xarray). Such custom index may also be used to hold some data that is tricky to propagate by other means, e.g., some internal information like “functional” coordinate parameters or a crs object. Xarray indexes should definitely provide more flexibility than coordinate data or attributes or accessor attributes for propagating this kind of information.

• Xarray accessors may be used to extend Dataset/DataArray public API. They could use the information stored in the CRSIndex/WCSIndex, e.g., add a crs read-only property that returns the crs object stored in CRSIndex, or add some some extract_crs_parameters method to extract the parameters and store them in Dataset/DataArray attributes similarly to what @djhoese suggests in his comment above.

For this use case a possible workflow would then be something like this:

# create or open an Xarray dataset with  x, y, z "pixel" (possibly lazy) coordinates
# and set a WCS index
dataset = (
    xr.Dataset(...)
    .set_index(['x', 'y', 'z'], WCSIndex, wcs_params={...})
)

# select data using pixel coordinates
dataset.sel(x=..., y=..., z=...)

# select data using world coordinates (via the "astro" accessor,
# which may access methods/attributes of the WCS index)
dataset.astro.sel_world(x=..., y=..., z=...)

# return a new dataset where the x,y,z "pixel" coordnates are replaced by the "world" coordinates
# (again using the WCS index, and propagating it to the returned dataset)
world_dataset = dataset.astro.pixel_to_world(['x', 'y', 'z'])

# select data using world coordinates
world_dataset.sel(x=..., y=..., z=...)

# select data using pixel coordinates (via the "astro" accessor)
world_dataset.astro.sel_pixel(x=..., y=..., z=...)

# this could be reverted
pixel_dataset = world_dataset.astro.world_to_pixel(['x', 'y', 'z'])
assert pixel_dataset.identical(dataset)

# depending on the implementation in WCSIndex, would either raise an error
# or implicitly convert to either pixel or world coordinates
xr.merge([world_dataset, another_pixel_dataset])

[benbovy]

@djhoese not yet but hopefully soon! Most of the work on explicit indexes is currently happening in #5692, which once merged (probably after the next release) will provide all the infrastructure for custom indexes. This is quite a big internal refactoring (bigger than I initially thought) that we cannot avoid as we’re changing Xarray’s core data model. After that, we’ll need to update some public API (Xarray object constructors, .set_index(), etc.) so that Xarray will accept custom index classes. This should take much less work than #5692, though.