Upsampling support

[pattonw]

I was about to try implementing a lazy downsampling of a dask array, but the downscale_dask looks perfect for my usecase.
I also need to lazily upsample a dask array. Would that be a useful addition to this library or should I keep that separate?
I was thinking I would just need to implement upscale_dask, upscale_coords, and upsampling_height and add them to the multiscale module.

[d-v-b]

yeah that sounds like a nice addition!

[pattonw]

I started down this path a little bit.
It seemed like upsampling might not fit so well as a direct inversion of the downsampling. Upscaling requires some context for the mapping so defining a WindowedInterpolator protocol and applying it to every key in the dask array doesn't seem like it would work.

I took a stab at using the high level dask interface:
Let me know if you have any thoughts on how I should improve it (If there are any downsides of using map_overlap that I don't know about), my knowledge of dask is still pretty high level at the moment.

# This is the block overlap. I'm only using 1 pixel overlap for now
# but this should probably depend on the interpolation order and upsampling factor
depth = [int(x > 1) for x in upsample]

# How much to slice off each block after having performed the upsampling.
# Equal to the amount of overlap multiplied by the scaling factor per axis.
trim_slicing = tuple(
    slice(d * s, (-d * s)) if d > 1 else slice(None)
    for d, s in zip(depth, upsample)
)


rescaled_arr = da.map_overlap(
    lambda x: rescale(
        x, upsample, order=1, preserve_range=True
    )[trim_slicing],  # upsample then trim
    array.data,
    depth=depth,
    boundary="reflect",  # must have some boundary padding so boundary blocks maintain shape after trimming
    trim=False,  # Cannot use built in trimming since it wouldn't account for upsampled boundary voxels
    dtype=array.data.dtype,
    chunks=tuple(c * u for c, u in zip(array.data.chunksize, upsample)),  # the shape of each chunk after upsampling
)

This seems to work but have some problems. I've noticed memory costs seem to rise with repeated random reads of large arrays that have been upscaled. I want the upscaling to be computed on the fly at each read since I am usually working with too much data to hold in memory but need to randomly sample many times and will eventually have read in the whole dataset.

[d-v-b]

although i don't have a lot of experience with it myself, I have heard that da.map_overlap can be pretty inefficient, and I know of at least one instance where the response to this was to manually implement the map_overlap logic directly. (cc @GFleishman)

If I understand things correctly, map_overlap in this case works by representing each input chunk as a separate dask task, and then converging those tasks on the worker that will execute the rescaling function. But this introduces two kinds of overhead -- dask has to schedule the tasks for the boundary chunks (and if we are in 3D, that's 8-fold more tasks than just those associated with the array itself), and dask may also optimize the task graph by trying to load boundary chunks on separate workers and sending the chunk data around via the network. This is probably inefficient.

I think a more efficient approach might be to keep dask in the dark about the central chunk / boundary chunk thing, and just map a function over the output chunks, where that function reads the (padded) data it needs to read from the source array, does the upsampling, and returns the upsampled data. If this is hard to express in xarray-multiscale, we might have to adjust this library a bit :)

[GFleishman]

+1 to @d-v-b's description, this is precisely my issue with da.map_overlap, though it's a bit worse than that. In 3D, each central chunk actually requires data from 26 neighbors to satisfy the overlap, don't forget about the corners. For each overlap, da.map_overlap spawns many tasks including overlap read, serialization, sending, and de-serialization; though I believe there are other tasks as well. I distribute very lengthy algorithms over large arrays, so I use a small compute chunk size. For an array size of 4096^3 and a central chunk size of 256^3, we have 4096 compute chunks. For something like that I would notice millions of tasks added to the task graph by da.map_overlap. The distributed scheduler has about a 1ms overhead per task, so 1M tasks results in over 15 minutes of time just to optimize the task graph.