Flexible coordinate transform

[benbovy]

This PR is a first step towards adding generic support for coordinate transforms in Xarray (i.e., analytical coordinates, functional index, etc.), which has been discussed already in different issues or threads:

• Rotation Functional Index example #8004
• Regular (linspace) Coordinates/Index #8473
• How can I implement a simple "Range Index" using custom Indexes? #8955
• My thoughts on coordinate zarr-developers/geozarr-spec#48
• https://discourse.pangeo.io/t/example-which-highlights-the-limitations-of-netcdf-style-coordinates-for-large-geospatial-rasters/4140

(I might miss other issues / discussions)

I started with a few rough experimentations but ended up with something more concrete that seems to work reasonably well, hence directly opening a (draft) PR. The design & implementation detailled below is still very much open to discussion, though! There's an usage example further below using a 2D affine transformation. It would be nice to test this with other examples.

cc @rabernat@dcherian@TomNicholas@martindurant

Design / Implementation

This PR adds three new classes that should facilitate integrating any coordinate transform into Xarray:

CoordinateTransform

Abstract (wrapper) class to handle coordinate transformation with support of dimension and coordinate names.

• many transforms should be pluggable via this class (by subclassing it)
• supports bulk (vectorized) transformation, both in forward and reverse direction
• supports any arbitrary number of coordinates / dimensions
  • lon, lat = f(x, y)
  • lon, lat, time = f2(x, y, t)
  • etc.
• one restriction is that the coordinates of a same transform must all have the same dimensions
  • lon(x,y) / lat(x,y)
  • lon(x,y,t) / lat(x,y,t) / time(x,y,t)
  • etc.
  • much simpler!
  • In some cases however, the transform parameter values are such that it can be applied independently over each dimension (e.g., rioxarray creates x(x) / y(y) dimension coordinates when the affine transform is rectilinear with no rotation). For those cases we cannot use a single CoordinateTransform instance, but it is still possible to wrap the same underlying transform object in several instances and link their respective coordinates at the xarray Index level (see below).

CoordinateTransformIndexingAdapter

Internal class for creating indexable coordinate variables from a transform (no need to change the Xarray data model!).

• wraps a CoordinateTransform instance
• coordinate labels are computed on-demand (lazy coordinates)
• supports both (explicit) orthogonal and vectorized indexing
• supports dimension re-ordering (transpose)
• doesn't support item assignment (of course)

CoordinateTransformIndex

Helper class for creating Xarray (custom) indexes based on coordinate transforms.

• wraps a CoordinateTransform instance
• takes care of creating the index (lazy) coordinates
• supports label-based selection (i.e., using "physical" or "world" labels)
  • only advanced (point-wise) indexing for now
  • any idea on what else would be nice here and how to implement it?
• supports alignment by comparing indexes based on their transform (not on their explicit coordinate labels)
  • only exact alignment for now (no join)
• may be used directly, although should mostly be either subclassed or encapsulated in another Xarray Index class
  • in the rioxarray example (see above), we might want to encapsulate two instances of CoordinateTransformIndex into a custom RasterIndex for the x and y coordinates respectively
  • a custom Xarray index is the right place for encoding / decoding coordinate definition

Usage Example (Affine 2D)

CoordinateTransform subclass

Let's write a subclass of CoordinateTransform that handles 2-d affine transformation (using affine). It is basically boilerplate code that takes care of dimension or coordinate names around the unlabelled input/output arrays of the underlying affine.Affine object.

importaffineimportxarrayasxrclassAffine2DCoordinateTransform(xr.CoordinateTransform):
"""Affine 2D coordinate transform."""affine: affine.Affinexy_dims=tuple[str]
def__init__(
self,
affine: affine.Affine,
coord_names: Iterable[Hashable],
dim_size: Mapping[str, int],
dtype: Any=np.dtype(np.float64),
):
# two dimensionsassertlen(coord_names) ==2assertlen(dim_size) ==2super().__init__(coord_names, dim_size, dtype=dtype)
self.affine=affine# array dimensions in reverse order (y = rows, x = cols)self.xy_dims=tuple(self.dims)
self.dims= (self.dims[1], self.dims[0])
defforward(self, dim_positions):
positions= [dim_positions[dim] fordiminself.xy_dims]
x_labels, y_labels=self.affine*tuple(positions)
results= {}
forname, labelsinzip(self.coord_names, [x_labels, y_labels]):
results[name] =labelsreturnresultsdefreverse(self, coord_labels):
labels= [coord_labels[name] fornameinself.coord_names]
x_positions, y_positions=~self.affine*tuple(labels)
results= {}
fordim, positionsinzip(self.xy_dims, [x_positions, y_positions]):
results[dim] =positionsreturnresultsdefequals(self, other):
returnself.affine==other.affineandself.dim_size==other.dim_size

Dataset, coordinates and index creation

In this example the index and the lazy coordinates are created from scratch, no pre-existing (explicit) coordinates are required!

fromxarray.indexesimportCoordinateTransformIndextransform=Affine2DCoordinateTransform(
affine.Affine.scale(1.0, 2.0),
coord_names=("xc", "yc"),
dim_size={"x": 10_000, "y": 20_000},
)
index=CoordinateTransformIndex(transform)
ds=xr.Dataset(coords=xr.Coordinates.from_xindex(index))

The resulting Dataset:

>>>ds<xarray.Dataset>Size: 3GBDimensions: (y: 10000, x: 20000)
Coordinates:
*xc (y, x) float642GB0.01.02.03.04.0 ... 2e+042e+042e+042e+04*yc (y, x) float642GB0.02.04.06.0 ... 1.999e+042e+042e+04Dimensionswithoutcoordinates: y, xDatavariables:
*empty*Indexes:
┌ xcCoordinateTransformIndex
└ yc

Coordinates "xc" and "yc" are big but they are lazy!

>>>ds.xc<xarray.DataArray'xc' (y: 10000, x: 20000)>Size: 2GB
[200000000valueswithdtype=float64]
Coordinates:
*xc (y, x) float642GB0.01.02.03.04.0 ... 2e+042e+042e+042e+04*yc (y, x) float642GB0.02.04.06.0 ... 1.999e+042e+042e+04Dimensionswithoutcoordinates: y, xIndexes:
┌ xcCoordinateTransformIndex
└ yc>>>ds["xc"].variable._dataCoordinateTransformIndexingAdapter(transform=<__main__.Affine2DCoordinateTransformobjectat0x15fb63790>)

Indexing

Orthogonal indexing (it is fast, it only computes 2x6 coordinate values below):

>>>ds.yc.isel(y=[0, 1, 3], x=slice(0, 2))
<xarray.DataArray'yc' (y: 3, x: 2)>Size: 48Barray([[0., 0.],
[2., 2.],
[6., 6.]])
Coordinates:
xc (y, x) float6448B0.01.00.01.00.01.0yc (y, x) float6448B0.00.02.02.06.06.0Dimensionswithoutcoordinates: y, x

Also works after re-ordering the dimensions:

>>>ds.transpose().yc.isel(y=[0, 1, 3], x=slice(0, 2))
<xarray.DataArray'yc' (x: 2, y: 3)>Size: 48Barray([[0., 2., 6.],
[0., 2., 6.]])
Coordinates:
xc (x, y) float6448B0.00.00.01.01.01.0yc (x, y) float6448B0.02.06.00.02.06.0Dimensionswithoutcoordinates: x, y

Vectorized indexing:

>>>ds.yc.isel(
... y=xr.Variable("points", [0, 1, 3]),
... x=xr.Variable("points", [0, 1, 3]),
... )
<xarray.DataArray'yc' (points: 3)>Size: 24Barray([0., 2., 6.])
Coordinates:
xc (points) float6424B0.01.03.0yc (points) float6424B0.02.06.0Dimensionswithoutcoordinates: points

Label-based selection

Point-wise selection:

>>>ds.sel(
... xc=xr.Variable("points", [101.34, 545.23, 876.76]),
... yc=xr.Variable("points", [13.12, 54.98, 76.43]),
... method="nearest",
... )
<xarray.Dataset>Size: 48BDimensions: (points: 3)
Coordinates:
xc (points) float6424B101.0545.0877.0yc (points) float6424B14.054.076.0Dimensionswithoutcoordinates: pointsDatavariables:
*empty*

What's next?

A few potential improvements from here:

• allow returning or re-calculating the transform instead of computing the coordinate labels while indexing
  • when possible, this should keep the xarray coordinates lazy and should also preserve their index
  • the obvious case if when a full slice is given for each dimension... Are there other less obvious cases?
  • allow CoordinateTransform.forward() to return a new instance of CoordinateTransform?
• allow special handling of dimension reduction (e.g., return another transform in the reduced space)
• possible to add generic support for joining / concatenating coordinate transforms? I.e., implement CoordinateTransformIndex.concat and CoordinateTransformIndex.join
• handle chunking
  • maybe best solved at a higher level? I.e., one transform instance per chunk
• add some convenient API for setting a new Xarray index from existing dimensions in a Dataset or DataArray?

[mdsumner]

Nice!! Thanks, I'm having fun with this - appreciate all the detail and functionality here it really helps a (non-native) Python learner.

[mdsumner]

One thing is that the coordinate values are currently "left"/"top" aligned, not the centre, so here we start at left/top 0,0 and end at 4,4.

fromxarray.indexesimportCoordinateTransformIndextransform=Affine2DCoordinateTransform(
affine.Affine.scale(1.0, 1.0),
coord_names=("xc", "yc"),
dim_size={"x": 5, "y": 5},
)
index=CoordinateTransformIndex(transform)
ds=xr.Dataset(coords=index.create_coordinates())
ds.yc.values#array([[0., 0., 0., 0., 0.],# [1., 1., 1., 1., 1.],# [2., 2., 2., 2., 2.],# [3., 3., 3., 3., 3.],# [4., 4., 4., 4., 4.]])

I'm assuming that a pure-scale puts us in the cell-area context of 0, 0, shape[0], shape[1] and so I pursued a world-realistic-ish context to convince myself.

I prefer to think in shape+bbox than in transforms when no shear is needed, so I'm using the gdal transform as an intermediate with a helper fun:

https://gist.github.com/mdsumner/dde0b611a4523e3485006c0df0143c2d

(fwiw, I'm sure this is obvious and not exactly priority rn but I'm excited to be able to delve into this and flesh out how I think about it in this context)

edit: I appreciate there's no absolutely right answer here, you might want (even decoupled per dimension) different alignment for your lazy coords in different contexts.

[benbovy]

Thanks for the feedback @mdsumner.

Yes the idea is to have something very generic in Xarray such that we can build domain-specific applications on top of it. It is very useful to test this functionality in various contexts now to make sure we are providing the right levels of abstractions. So please keep having fun with this :-) !

Regarding your example, I think that rioxarray combines the input affine transformation with affine.Affine.translation(0.5, 0.5) to make coordinate values center aligned. If it is more natural to think in shape+bbox than in transforms in the geo domain, let's build something on top of CoordinateTransformIndex, e.g., something like below adapted from your gist example:

classGeoIndex(CoordinateTransformIndex):
@classmethoddeffrom_shape(cls, shape, bbox=None, center=True):
ifbboxisNone: bbox= (0.0, 0.0, shape[0], shape[1])
gdal= (
bbox[0], (bbox[2] -bbox[0]) /shape[0], 0.0, bbox[3], 0.0, (bbox[1] -bbox[3]) /shape[1]
)
aff=affine.Affine.from_gdal(*gdal)
ifcenter:
coord_names= ("xc", "yc")
aff*=affine.Affine.translation(0.5, 0.5)
else:
# left/topcoord_names= ("xlt", "ylt")
transform=Affine2DCoordinateTransform(
aff,
coord_names=coord_names,
dim_size={"x": shape[1], "y": shape[0]},
)
returncls(transform=transform)

>>>bbox= (-3950000, -3950000, 3950000, 4350000)
>>>shape= (316, 332)
>>>index=GeoIndex.from_shape(shape, bbox=bbox)
>>>ds=xr.Dataset(coords=xr.Coordinates.from_xindex(index))
>>>ds.isel(x=slice(0, 159), y=slice(0, 167))
<xarray.Dataset>Size: 425kBDimensions: (y: 167, x: 159)
Coordinates:
xc (y, x) float64212kB-3.938e+06-3.912e+06 ... -1.25e+041.25e+04yc (y, x) float64212kB4.338e+064.338e+06 ... 1.875e+051.875e+05Dimensionswithoutcoordinates: y, xDatavariables:
*empty*

(This gives the same coordinate values than the "ice" dataset loaded from the .tif file using the rasterio engine in your 2nd gist)

[benbovy]

This is important I think.

If the coordinates values correspond to the physical values at the top/left pixel corners in the 2D case, we may rather want np.floor(pos).astype("int") when converting decimal positions (obtained by inverse transformation) to integer indexers.

[martindurant]

Great to see this. I haven't looked at the implementation yet, but I think I agree with the description whole heartedly.

It would be the place of the various IO backends to instantiate the affine (or whatever) transform from the the metadata standards of the respective formats.

[benbovy]

For completeness, here is an implementation of the 1-dimensional "range index" discussed in #8955.

The coordinate transform subclass:

classRange1DCoordinateTransform(xr.CoordinateTransform):
"""Simple bounded interval 1-d coordinate transform."""left: floatright: floatdim: strsize: intdef__init__(
self,
left: float,
right: float,
coord_name: Hashable,
dim: str,
size: int,
dtype: Any=None,
): ifdtypeisNone:
dtype=np.dtype(np.float64)
super().__init__([coord_name], {dim: size}, dtype=dtype)
self.left=leftself.right=rightself.dim=dimself.size=sizedefforward(self, dim_positions):
positions=dim_positions[self.dim]
labels=self.left+positions* (self.right-self.left) /self.sizereturn {self.dim: labels}
defreverse(self, coord_labels):
labels=coord_labels[self.coord_names[0]]
positions= (labels-self.left) *self.size/ (self.right-self.left)
return {self.dim: positions}
defequals(self, other):
return (
self.left==other.leftandself.right==other.rightandself.size==other.size
)

Dataset creation:

>>>range_tr=Range1DCoordinateTransform(1.0, 2.0, "x", "x", 100)
>>>index=CoordinateTransformIndex(range_tr)
>>>ds=xr.Dataset(data_vars={"foo": ("x", np.arange(100))}, coords=xr.Coordinates.from_xindex(index))
>>>ds<xarray.Dataset>Size: 2kBDimensions: (x: 100)
Coordinates:
*x (x) float64800B1.01.011.021.031.04 ... 1.961.971.981.99Datavariables:
foo (x) int64800B0123456789 ... 919293949596979899Indexes:
xCoordinateTransformIndex

This example is interesting because in this simple case we would expect a few more operations to work than in the case of more complex transformations such as 2D affine with rotation and/or shear, e.g.,

• indexing with a slice (step=1) should preserve the coordinate index but it doesn't:

>>>ds.isel(x=slice(5, 10)).xindexesIndexes:
*empty*

• basic label-based selection should also work, but it is not supported:

>>>ds.sel(x=1.65, method="nearest")
TypeError: CoordinateTransformIndexonlysupportsadvanced (point-wise) indexingwitheitherxarray.DataArrayorxarray.Variableobjects.

Perhaps we could try adding support for this in CoordinateTransform and/or CoordinateTransformIndex? My concern is that we may end up cluttering the interface / implementation of those classes with many special cases.

An alternative option is building on top of it, e.g., in this case also provide a Range1DIndex class like so:

---- expand here to see the implementation of Range1DIndex ----

fromxarray.core.indexesimportIndexSelResultclassRange1DIndex(CoordinateTransformIndex):
transform: Range1DCoordinateTransformdim: strcoord_name: Hashablesize: intdef__init__(
self,
left: float,
right: float,
coord_name: Hashable,
dim: str,
size: int,
dtype: Any=None,
):
self.transform=Range1DCoordinateTransform(
left, right, coord_name, dim, size, dtype
)
self.dim=dimself.coord_name=coord_nameself.size=sizedefisel(self, indexers):
idxer=indexers[self.dim]
# straightforward to generate a new index if a slice is given with step 1ifisinstance(idxer, slice) and (idxer.step==1oridxer.stepisNone):
start=max(idxer.start, 0)
stop=min(idxer.stop, self.size)
new_left=self.transform.forward({self.dim: start})[self.coord_name]
new_right=self.transform.forward({self.dim: stop})[self.coord_name]
new_size=stop-startreturnRange1DIndex(new_left, new_right, self.coord_name, self.dim, new_size)
returnNonedefsel(self, labels, method=None, tolerance=None):
label=labels[self.dim]
ifisinstance(label, slice):
iflabel.stepisNone:
# slice indexing (preserve the index)pos=self.transform.reverse({self.dim: np.array([label.start, label.stop])})
pos=np.round(pos[self.coord_name]).astype("int")
new_start=max(pos[0], 0)
new_stop=min(pos[1], self.size)
returnIndexSelResult({self.dim: slice(new_start, new_stop)})
else:
# otherwise convert to basic (array) indexinglabel=np.arange(label.start, label.stop, label.step)
# support basic indexing (in the 1D case basic vs. vectorized indexing# are pretty much similar)unwrap_xr=Falseifnotisinstance(label, xr.Variable|xr.DataArray):
# basic indexing -> either scalar or 1-d arraytry:
var=xr.Variable("_", label)
exceptValueError:
var=xr.Variable((), label)
labels= {self.dim: var}
unwrap_xr=Trueresult=super().sel(labels, method=method, tolerance=tolerance)
ifunwrap_xr:
dim_indexers= {self.dim: result.dim_indexers[self.dim].values}
result=IndexSelResult(dim_indexers)
returnresult

>>>index=Range1DIndex(1.0, 2.0, "x", "x", 100)
>>>ds2=xr.Dataset(data_vars={"foo": ("x", np.arange(100))}, coords=xr.Coordinates.from_xindex(index))

Slicing (notice the preserved Range1DIndex):

>>>ds2.isel(x=slice(5, 10))
<xarray.Dataset>Size: 80BDimensions: (x: 5)
Coordinates:
*x (x) float6440B1.051.061.071.081.09Datavariables:
foo (x) int6440B56789Indexes:
xRange1DIndex

Some basic label-based selection:

>>>ds2.sel(x=1.654, method="nearest")
<xarray.Dataset>Size: 16BDimensions: ()
Coordinates:
xfloat648B1.65Datavariables:
fooint648B65>>>ds2.sel(x=slice(1.465, 1.874), method="nearest") # preserves the index!<xarray.Dataset>Size: 640BDimensions: (x: 40)
Coordinates:
*x (x) float64320B1.471.481.491.51.51 ... 1.831.841.851.86Datavariables:
foo (x) int64320B4748495051525354 ... 7980818283848586Indexes:
xRange1DIndex

For such a simple 1-d range example, the coordinate transform abstraction is actually a bit overkill but still has the advantage of providing the lazy coordinate variable "for free".

[astrofrog]

@benbovy - @Cadair and I have been playing around with trying to get this to work with the astropy APE 14 WCS specification. Here is a minimal example:

https://gist.github.com/Cadair/4a03750868e044ac4bdd6f3a04ed7abc

We are running into a bug in the __repr__ which is causing an out of bounds error. It seems that accessing the coordinates directly works so it seems to be a problem specific to the __repr__?

Another unrelated comment: it would be nice to have the CoordinateTransform class be a proper abc class, and have the methods that need to be implemented be defined as abstract methods (e.g. forward and reverse)

[benbovy]

Thanks for the feedback @astrofrog!

I'll look into the __repr__ issue. Could you provide a minimal reproducible example or a link where I can download the FITS file used in your example, please?

[Cadair]

@benbovy the fits file is included with astropy , so the code in the notebook should run as-is I believe.

[benbovy]

Ah thanks. The __repr__ issue should now be fixed in 09667c5.

[shoyer]

How hard would it be to support tolerance in some form? This is a common and useful form of error checking.

[benbovy]

Pretty tricky to support it here I think, probably better to handle it on a per case basis.

For basic transformations I guess it could be possible to calculate a single, uniform tolerance value in decimal array index units and validate the selected elements using those units (cheap). In other cases we would need to compute the forward transformation of the extracted array indices and then validate the selected elements based on distances in physical units (more expensive).

Also, there may be cases where the coordinates of a same transform object don’t have all the same physical units (e.g., both degrees and radians coordinates in an Astropy WCS object). Unless we forbid that in xarray.CoordinateTransform, it doesn’t make much sense to pass a single tolerance value. Passing a dictionary tolerance={coord_name: value} doesn’t look very nice either IMO. A {unit: value} dict looks better but adding explicit support for units here might be opening a can of worms.

[shoyer]

This very exciting! Nice work.

For indexing, it may be worth considering if you can implement .interp(). In practice I think that is often more desirable than nearest neighbor lookup.

[rbavery]

rioxarray creates x(x) / y(y) dimension coordinates when the affine transform is rectilinear with no rotation). For those cases we cannot use a single CoordinateTransform instance

I think rioxarray does this for performance reasons because it is faster and possible to correctly apply the affine transformation without calling numpy.meshgrid when the affine is rectilinear with no rotation. But with flexible coordinates, I think both approaches could be replaced with only a single CoordinateTransform with some refactoring.

importnumpyimportaffinetransform=affine.Affine.translation(.5,.5).scale(1.0)
width=512height=200

%%timex_coords, _=transform* (numpy.arange(width), numpy.zeros(width))
_, y_coords=transform* (numpy.zeros(height), numpy.arange(height))

Wall time: 290 μs

%%timex_coords_mesh, y_coords_mesh=transform*numpy.meshgrid(
numpy.arange(width),
numpy.arange(height),
)

Wall time: 3.09 ms

possible to add generic support for joining / concatenating coordinate transforms? I.e., implement CoordinateTransformIndex.concat and CoordinateTransformIndex.join

This sounds very valuable but want to make sure I understand what is meant. If I have a dataset of rasters across different UTM projections, would this allow me to read each with rioxarray and then concatenate the raster arrays such that each raster maintains it's original CRS? Or would this enable concatenating rasters that are already in the same CRS? Or something else?

My use case for this is I'd like to avoid reprojection and have a single xarray.DataArray representing rasters spread over global extents. And I'd like to be able to save this concatenated xarray DataArray to a Zarr v3 store with sharding in a way that preserves each CRS, with GeoZarr.

[benbovy]

both approaches could be replaced with only a single CoordinateTransform with some refactoring.

Hmm do you have an idea on how this refactoring would look like? I've tried implementing a version of CoordinateTransform that supports coordinates with different dimensions but I eventually gave up because it was too complicated.

Here is one way to support the rectilinear / no rotation affine transform with independent x, y 1-dimensional coordinates without any refactoring:

• a CoordinateTransform subclass that wraps an affine.Affine instance for either the x or y coordinate

---- expand here to see the implementation of AxisAffineCoordinateTransform ----

classAxisAffineCoordinateTransform(xr.CoordinateTransform):
"""1-axis wrapper of an affine 2D coordinate transform with no skew/rotation. """affine: affine.Affineis_xaxis: boolcoord_name: Hashabledim: strsize: intdef__init__(
self,
affine: affine.Affine,
coord_name: Hashable,
dim: str,
size: int,
is_xaxis: bool,
dtype: Any=np.dtype(np.float64),
):
if (notaffine.is_rectilinearor (affine.b==affine.d!=0)):
raiseValueError("affine must be rectilinear with no rotation")
super().__init__((coord_name,), {dim: size}, dtype=dtype)
self.affine=affineself.is_xaxis=is_xaxisself.coord_name=coord_nameself.dim=dimself.size=sizedefforward(self, dim_positions):
positions=dim_positions[self.dim]
ifself.is_xaxis:
labels, _=self.affine* (positions, np.zeros_like(positions))
else:
_, labels=self.affine* (np.zeros_like(positions), positions)
return {self.coord_name: labels}
defreverse(self, coord_labels):
labels=coord_labels[self.coord_name]
ifself.is_xaxis:
positions, _=~self.affine* (labels, np.zeros_like(labels))
else:
_, positions=~self.affine* (np.zeros_like(labels), labels)
return {self.dim: positions}
defequals(self, other):
returnself.affine==other.affineandself.dim_size==other.dim_size

• an Xarray Index that encapsulates two CoordinateTransformIndex instances (sharing the same Affine object) for the x and y axis respectively

---- expand here to see the implementation of RasterIndex ----

fromxarrayimportVariablefromxarray.indexesimportCoordinateTransformIndexfromxarray.core.indexingimportIndexSelResult, merge_sel_resultsclassRasterIndex(xr.indexes.Index):
def__init__(
self,
x_index: CoordinateTransformIndex,
y_index: CoordinateTransformIndex,
):
self.x_index=x_indexself.y_index=y_index@classmethoddeffrom_transform(
cls,
affine: affine.Affine,
shape: tuple[int, int],
xy_coord_names: tuple[Hashable, Hashable] = ("x", "y"),
):
# shape is in y, x orderxtr=AxisAffineCoordinateTransform(
affine, xy_coord_names[0], xy_coord_names[0], shape[1], is_xaxis=True
)
ytr=AxisAffineCoordinateTransform(
affine, xy_coord_names[1], xy_coord_names[1], shape[0], is_xaxis=False
)
returncls(CoordinateTransformIndex(xtr), CoordinateTransformIndex(ytr))
defcreate_variables(
self, variables: Mapping[Any, Variable] |None=None
) ->dict[Hashable, Variable]:
return {**self.x_index.create_variables(), **self.y_index.create_variables()}
defcreate_coords(self) ->xr.Coordinates:
variables=self.create_variables()
indexes= {name: selffornameinvariables}
returnxr.Coordinates(coords=variables, indexes=indexes)
defsel(
self, labels: dict[Any, Any], method=None, tolerance=None
) ->IndexSelResult:
results= []
xlabels= {k: vfork, vinlabelsifkinself.x_index.transform.coord_names}
ifxlabels:
results.append(self.x_index.sel(xlabels))
ylabels= {k: vfork, vinlabelsifkinself.y_index.transform.coord_names}
ifylabels:
results.append(self.y_index.sel(ylabels))
returnmerge_sel_results(results)
defequals(self, other: Self) ->bool:
returnself.x_index.equals(other.x_index) andself.y_index.equals(other.y_index)

Usage example:

>>>index=RasterIndex.from_transform(affine.Affine.translation(0.5, 0.5), (1000, 2000))
>>>ds=xr.Dataset(coords=xr.Coordinates.from_xindex(index))
>>>ds<xarray.Dataset>Size: 24kBDimensions: (x: 2000, y: 1000)
Coordinates:
*x (x) float6416kB0.51.52.53.5 ... 1.998e+031.998e+032e+03*y (y) float648kB0.51.52.53.54.5 ... 996.5997.5998.5999.5Datavariables:
*empty*Indexes:
┌ xRasterIndex
└ y>>>ds.isel(x=slice(100, 200), y=500)
<xarray.Dataset>Size: 808BDimensions: (x: 100)
Coordinates:
x (x) float64800B100.5101.5102.5103.5 ... 197.5198.5199.5yfloat648B500.5Datavariables:
*empty*

[benbovy]

My use case for this is I'd like to avoid reprojection and have a single xarray.DataArray representing rasters spread over global extents. And I'd like to be able to save this concatenated xarray DataArray to a Zarr v3 store with sharding in a way that preserves each CRS, with GeoZarr.

This seems complicated to me. It would be easier in this case to have a unique CRS per DataArray and provide a virtual layer built on top of DataArray to handle lazy reprojection (e.g., similarly to GDAL VRT I guess?).

Also, IIUC Zarr doesn't allow per-chunk or per-shard metadata so it isn't clear to me how GeoZarr would support multiple CRS datasets (zarr-developers/geozarr-spec#4).

[martindurant]

Zarr doesn't allow per-chunk or per-shard metadata

perhaps virtualizarr can do this (@TomNicholas )

[benbovy]

@dcherian Yes let's do this (unless we want to avoid breaking changes in the next related PRs... although I'm already pretty happy with the API here). I'll add unit tests in a follow-up PR.

[keewis]

we want to avoid breaking changes in the next related PRs

I think the idea was to merge but not expose as public API for now, so that we don't need to worry about breaking changes

[benbovy]

Hmm maybe we should fix CI before merging it, though? I can have a look tomorrow.

[dcherian]

yes exactly, just fix mypy, merge and don't tell anyone about it till we're ready haha. Let's clearly add an experimental API warning to the docstring too.

[maxrjones]

Thanks for this PR - super excited to try it out!

These suggested changes make mypy happy again.

[benbovy]

I added the tests here (mostly complete for the features added). This should be ready for review (or for merging :-)).

[dcherian]

I took a quick look through the tests. Thanks @benbovy! This is a large leap forward.