Explore Xarray as the basis for a genetic toolkit API

[eric-czech]

I'm not sure what advantages having labeled axes has for big call matrices, but xarray may make sense as a way to carry along the variant and sample metadata (as opposed to dask dfs). It will be worth a shot to see how useful that interface is.

[eric-czech]

Notable Xarray limitations found so far are:

• There is no support for SQL joins
  • It is possible to do a Hail-style "join", but this is really just index alignment (https://github.com/pydata/xarray/issues/3791)
  • I'll have to check, but I think index alignment may only work for a univariate index and not a multiindex (https://github.com/pydata/xarray/issues/1603)
  • This would mean we'd have to combine fields like chromosome and position before indexing+joining
• Coordinates need to fit in memory (https://github.com/pydata/xarray/issues/1650)
  • Even if data variables represent things like say contig, position, and alleles (a common primary key for variants), running an index alignment on external data for say rsID will require conversion to coordinates
• Masking is not supported (https://github.com/pydata/xarray/issues/1194). It is supported in Dask (dask/masked-arrays) but not in CuPy (https://github.com/cupy/cupy/issues/2225)
• No map_overlap support (https://github.com/pydata/xarray/issues/3147)
• Chunk sizes for underlying dask arrays need to be uniform (https://github.com/pydata/xarray/issues/2225)
  • This could be important for having chunks broken on contig, which would simply methods like LD pruning
• Dataset/DataArray polymorphism is not supported (https://github.com/pydata/xarray/issues/706)
  • The recommendation is to use accessors but these are attached globally to every single Dataset instance, which is awkward for trying to support multiple subclasses
• Neither nominal nor structural subtype safety (via static analysis) is possible (https://github.com/pydata/xarray/issues/3959)

Note: some of these were aggregated from https://github.com/related-sciences/rs-platform/issues/19#issuecomment-594211481

[eric-czech]

There isn't much remaining doubt about Xarray. Updates on some of the limitations above:

• masked data: Whether we use masking or NA sentinels, this is a numpy problem so it needs to be solved at that level (so any PyData solution will have this issue)
• nominal/structural type safety: This also needs to be solved on the numpy level
• Dataset/DataArray subclassing: I couldn't make Xarray accessors useful in the way I wanted without monkey-patching, but they work well enough and there may be a cleaner way to do what I did
• map_overlap support: This is in Dask now at least so it wouldn't be that hard to introduce in Xarray or otherwise work around
• uniform chunk sizes: While this is true when a dataset needs to be written as zarr, rechunking before write may not be as big of performance issue as I thought (I've found that its effects are generally marginal compared to everything else)

The others are still legitimate limitations but don't affect core GWAS operations.