Discussion : writing xarray-compatible zarr stores

[emileten]

Context

Hi @joemoorhouse. This topic was briefly mentioned in this issue os-climate/physrisk#263 when we discussed how coordinate labels are stored in, or retrieved from, the output zarr stores in OS-C, in particular the hazard indicators. It was stressed that we do not save coordinate labels directly in the zarr metadata and instead, save a matrix that can be used in combination with the indices to get the coordinates. This makes the output incompatible with native xarray operations. This is why we have a module dedicated to offering an api to read the output as xarray objects

Suggestion

Can we write the coordinates to the zarr metadata so as to have the output directly xarray-compatible ? This would for example allow us to experiment with titiler-xarray for serving tiles..

You mentioned performance concerns in the linked issue, but I am not sure to understand why, in principle, this issue should arise, as the ideal way, I believe (and there is some reference to that), to save such data would be one single zarr store with latitude, longitude, temporal, model and scenario dimensions (and with the coordinate labels for each of those in the zarr store metadata), with a chunking structure optimized for the usage of this data. This also suggests potentially broader changes in how the data is saved, since I think, for example for the days above tas workflow, we're saving separate zarr arrays for each year ?

[sharkinsspatial]

@emileten There is a currently an additional utility write_data_array method on the OSCZarr class that supports writing with xarray compliant _ARRAY_DIMENSIONS convention. This works for me when used in the pytest environment but is throwing an xarray exception on set_index when run with actual nex_gddp_cmip6 data, I'll need to investigate a bit more here.

@joemoorhouse if possible, can you elaborate a bit on the use of an index variable in the normalize_array method. I'm not clear on the purpose for indexing on an additional variable here?

Additionally, looking through some of the discussion on the interpolation methods in the physrisk repo, are there selection/interpolation methods being used in physrisk that can't be supported using standard xarray coordinates and interpolation routines?

I may be misunderstanding the intent, but it seems like some of the code could be potentially simplified by working directly in xarray with the zarr data in transform space rather than numpy with data in raw pixel space. I know this was discussed a bit here but can you provide a bit more detail on the advantages of the "Geotiff" based approach being used.

[joemoorhouse]

Hi @emileten, hi @sharkinsspatial,
Some answers in reverse order (and sounds like a call to discuss might also be useful?):
First on the question of why in physrisk we are not using xarray (as compared to hazard where we use widely), but rather we are using the Zarr library directly to access the Zarr array. Indeed I started by using xarray in physrisk but was fighting to get the performance I wanted, so opted to go lower level: there is not muchxarray functionality we need and performance optimisation was harder. In hazard, in contrast, our use cases already seem a good fit to using xarray + Dask.
More precisely, we have a (very) important main use case which is looking up single points from latitudes and longitudes. There we want to look up ~ millions of locations as quickly as possible. Here the affine transform gives a computational advantage because we can calculate the pixel locations directly: no look-up necessary at all. When doing the same thing with xarray labels I guess (?) we are at best doing a log2(N) complexity search for each look up along each co-ordinate axis, where N is the dimension length (from memory: I'd have to debug through this again to check what it is doing!).
Other considerations too:

• We want to Zarr async behaviour in downloading chunks, but we need the flexibility to control any thread/process pool behaviour on top of that; I'm a bit wary about what xarray is doing there.
• Nice to avoid loading up and caching labels, e.g. in the API.

[joemoorhouse]

The index dimension is just a convention whereby the hazard indicators are 3 dimensional Zarr arrays by construction. Two dimensions are spatial and the third is the index dimension that depends on the nature of the data. Acute hazards are 3D because we always have a number of return periods e.g. flood depth for 30Y, 100Y, 200Y, ... returns. For some chronic hazards it is also useful to make 3D because we have data that depends on a number of thresholds, e.g. days with wet bulb globe temperature above 20, 25, 30, ... C.

It is important performance-wise to have something 3D, because for the physical climate risk calculations we typically need all of the index values for a particular lat and lon (so also want to have in the same chunk).

[joemoorhouse]

On storing the coordinate labels, I think when the xarray [to_zarr method] (https://docs.xarray.dev/en/latest/generated/xarray.Dataset.to_zarr.html) is called the xarray is actually stored as a Zarr group with arrays for the data itself and each coordinate. That is, the convention is not to store the labels in the meta-data but have them in separate arrays. @emileten, was that what you had in mind, or was it actually a variation on that?

I suspect that the to_zarr method format is trying to keep the metadata (attributes) as small as possible so that it is fast to load an array - although you cannot then create an xarray without downloading the coordinates. I was also wondering whether to adopt that convention...

I am in two minds. Perhaps personal taste but I find the GeoTiff-like approach of storing the affine transform much more elegant than storing all of the labels: that way we can store the labels in the metadata without this getting too big because we are storing the labels so efficiently (6 numbers).

I also struggle a bit with the structure we end up with when using the to_zarr method because an array becomes a group. I tried this out for representing map tile sets where we end up with a group of say 9 zoom levels and each zoom level is itself a group containing the array and coordinates. I thought it seemed an overly complex structure.

On the other hand, adopting that convention does not prevent us from doing any performance optimisations, i.e. we could still ignore the co-ordinates arrays and regenerate from affine transform if needed.

Whatever option we choose, it is still pretty easy to create an xarray from the Zarr array and it might not prevent us from using say titiler-xarray (e.g. if we can add a hook for adding the logic to create the xarray from a single Zarr array, as opposed to a group).

[emileten]

Thanks both for all the input. For now we'll just ensure this module offers the option to write output with coordinates in the metadata : #67. This will be optional. The default remains the legacy behavior.

[emileten]

Marked this as a 'discussion', it's not really an issue anymore. Not sure how to convert an issue to an actual github discussion.

[j08lue]

Related discussion on Pangeo considering Xarray to support affine transforms in gridded data stores (e.g. Zarr):

• https://discourse.pangeo.io/t/example-which-highlights-the-limitations-of-netcdf-style-coordinates-for-large-geospatial-rasters/4140/6

[joemoorhouse]

Hi @j08lue and FYI @xbarra,

Adding some of our exchanges into this discussion (can still discuss face to face as needed!). Some renewed focus because we have the question of 'what is a nice data store for ASDI'.

@j08lue, as you note the functionality to support affine transforms is, most likely, making its way into the XArray core library pydata/xarray#9543 and presumably at some point after that there will be an XArray-native way to save co-ordinates into Zarr as well as a GeoZarr standard. Separate discussion on that: zarr-developers/geozarr-spec#17. And thanks for the nice demo of how we can store data as a data cube, using the NetCDF-style coordinates of XArray:

https://gist.github.com/j08lue/193291492504d61ad9f446e8459380ef#file-os-climate-hazard-reformat-zarr-multidimensional-with-stac-ipynb

I think it is fair to say that at the point there is a GeoZarr standard we should align to that. If anything the discussion you mention above reinforces that a transform is preferable to NetCDF-style coordinates in our case: that's surely the best option for us long-term.
https://discourse.pangeo.io/t/example-which-highlights-the-limitations-of-netcdf-style-coordinates-for-large-geospatial-rasters/4140/6

I think in the meantime we have two questions.

1. Until the new standard is relatively clear and supported, should we also have NetCDF-style coordinates?
2. Should we have more of a data cube structure as the gist example shows and if so, how?

For 1 I am in two minds. Compared to today, having the NetCDF coordinates means that we have to add another level into the storage scheme I believe (which may or may not be redundant in the eventual target). That is in Zarr nomenclature we have a group 'days_tas_above' of arrays 'days_tas_above' (actual indicator data), 'latitude', 'longitude', 'temperature' (coordinates) etc. The arrays would then be at paths:
.../days_tas_above/days_tas_above, .../days_tas_above/latitude, ../days_tas_above/longitude, ../days_tas_above/temperature etc
We can still have the transform as an attribute in the .../days_tas_above/days_tas_above array, so in principle we can have optional NetCDF coordinates.
I guess my main concern is if adding this extra level takes us further away from the long-term structure, which might have the transform as attributes in the array itself, i.e. the real target might actually be closer to what we have now (I'm pretty sure the attributes will be completely different of course :)). But on the flip side, right now we cannot open an XArray without some extra logic. Also, I don't think the extra layer will actually limit out options in the future.

In an effort to conclude, I suggest that for now we support the cases both with and without NetCDF coordinates in physrisk, but we do not migrate one format to another (at least until clear what the long-term looks like)

2 is perhaps a bit easier. I agree 100% with the gist example that the temperatures should be in the array (and in fact we started doing just that in other cases). I also agree with that example that emissions scenario should be collection-level information. I guess I would probably make time (i.e. the projection central year) collection-level information too rather than a dimension. My thinking for that is probably driven by the use cases I have seen. There are cases where a user is interested just in one or two scenarios and one or two years and a data cube perhaps binds scenarios and times too closely together for my liking (i.e. in the same Zarr array structure even if these are in different chunks). For example, if one wanted to copy just the 2050 data that would become a Zarr/XArray operation rather than just copying using a different S3 prefix. Another case where I can see arguments both ways.