Meta-issue: summary of work on non-reprojected dimensions and improved dask support

[astrofrog]

@Cadair @svank - I have made a number of improvements/changes since v0.21.0 while trying to get some larger-that-memory mosaicking jobs working, and I thought it might be helpful if I summarize the changes here.

New options

reproject_interp, reproject_adaptive, reproject_exact

• return_type='zarr' and zarr_path= (Added intermediate_memmap='zarr' option which dumps to temporary zarr arrays #612): in addition to 'numpy' and 'dask', the output can now be written to a zarr store on disk at zarr_path (which must not already exist). The reprojection is then always computed in blocks (block_size defaults to 'auto' if not given) and dask arrays backed by the zarr store are returned. We were writing zarr arrays internally anyway, so this is just a way to shortcut things to avoid having to them convert zarr arrays to numpy in cases where zarr is good enough.
• non_reprojected_dims= (Make it so we can use broadcasting even if input WCS is same dimension as data #539; reproject_interp and reproject_adaptive only): allows broadcasting the reprojection over leading dimensions even when the input/output WCS have the same dimensionality as the data, assuming a one-to-one pixel mapping along those dimensions. The main use case is e.g. cubes where each slice (time, spectral) maps cleanly onto one output slice but the WCS is genuinely N-d (e.g. spatial drift over time). For now leading dimensions must be given as sequential integers from zero ((0,), (0, 1), ...), and currently requires a block_size matching shape_out along the reprojected dimensions (i.e. slice-by-slice processing, optionally with parallel=). In future we could allow arbitrary tuples if we wanted.

reproject_and_coadd

This is the biggest chunk of work (#388 plus follow-ups):

• return_type= ('numpy' | 'dask' | 'zarr', default 'numpy' = old behavior):

  • 'dask' (Implement option to use dask to do image co-addition #388, Assemble each output chunk of the dask coaddition from only the images that overlap it and reduce memory usage for big-endian data #616): builds the whole co-addition as a single deferred dask graph and returns uncomputed dask arrays, so e.g. a cutout of a mosaic can be computed without ever materializing the full mosaic. Each output chunk is assembled from only the images that overlap it (Assemble each output chunk of the dask coaddition from only the images that overlap it and reduce memory usage for big-endian data #616), and the combination matches the numpy path exactly (footprint-weighted mean/sum, footprint-aware first/last/min/max); input_weights are supported. Not compatible with match_background, output_array/output_footprint, or intermediate_memmap.
  • 'zarr' (Add support for return_type='zarr' in reproject_and_coadd #621): builds the same graphs but computes them immediately, batch by batch, into a zarr store at zarr_path, so peak memory is bounded by the batch size regardless of mosaic size. zarr_batch_size= controls how many output chunks are computed per batch (default targets ~2 GB per batch); batches are computed with the scheduler implied by parallel=, with the same semantics as the individual reprojection functions. I found this was needed otherwise with a plain return_type='dask', the graph could end up being complicated and dask sometimes computed it in a non-optimal way memory-wise

• non_reprojected_dims= (Add support for non_reprojected_dims in reproject_and_coadd #611): co-add over non-reprojected leading dimensions (see above), passed through to the reprojection function. Fix computation of footprint in co-addition when non-reprojected dimensions are present #615 fixed the tile footprint computation so this also works when the output WCS has fewer dimensions than the input, e.g. a celestial-only output WCS for a set of cubes.

• combine_function='median' is now supported for all return types: for 'dask'/'zarr' as part of the deferred graph (Implement option to use dask to do image co-addition #388), and for 'numpy' via a chunked combine (Implement support for chunked median with return_type='numpy' in reproject_and_coadd #622) that retains the reprojected arrays until the end, as match_background does (so intermediate_memmap can keep them on disk), which also makes the median composable with match_background. The median is unweighted.

• intermediate_memmap='zarr' (Added intermediate_memmap='zarr' option which dumps to temporary zarr arrays #612): store intermediate reprojected arrays as temporary zarr arrays instead of numpy memmaps (implemented via the new return_type='zarr' on the reprojection functions). Each image is reprojected in blocks and its store is removed as soon as it has been combined, so disk usage does not accumulate. Not compatible with match_background. The idea here was that since these were internal temporary files anyway, it was faster to just use zarr and not waste time converting to-from memory mapped arrays on disk (which has to be done serially)

• block_sizes can now be a single block size (tuple or 'auto') applied to all datasets, as well as the previous one-per-dataset list; with return_type='dask' a single common block size is also used as the chunking of the returned dask arrays. block_size= (singular) is accepted as a passthrough alias (Implement option to use dask to do image co-addition #388).

  Other changes

• The internal pixel_to_pixel_with_roundtrip helper was replaced by a new chunked pixel_to_pixel_chunked (Significantly reduce memory usage due to WCS transformations #617) - which avoids having memory blow up when using GWCS, and is actually a bit faster.

• There are also a number of bug fixes, performance improvements etc which I won't detail here.

My next step is going to be to document all this, but I am first going to try and reorganize the docs, which are a bit of a mess currently.