astropy / astropy/reproject

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

Open
#623 0 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
127
Forks
74
Avg merge
1d 11h
Merged PRs (30d)
2

Description

@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=`** (#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=`** (#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'` (#388, #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 (#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'` (#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=`** (#611): co-add over non-reprojected leading dimensions (see above), passed through to the reprojection function. #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 (#388), and for `'numpy'` via a chunked combine (#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'`** (#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 (#388).

### Other changes

* The internal `pixel_to_pixel_with_roundtrip` helper was replaced by a new chunked `pixel_to_pixel_chunked` (#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.

Contributor guide

No contributing guide indexed for this repository

Research direction

Review the existing documentation for reproject_interp, reproject_adaptive, reproject_exact, and reproject_and_coadd, then map the options and changes summarized here into the reorganized docs. Done means the new return types, non_reprojected_dims, median, intermediate storage, block sizing, and chunked pixel conversion behavior are documented.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
documentation
Issue type
Documentation
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.