I'm unsure if this PR is worth it, but it offers some extra analysis on the optimization/chunk size front.

This PR speeds up computation with the distributed scheduler a bit by changing the code for building the Dask array to simplify the task graph (or so I believe). Instead of several iterations of da.stack calls, it directly generates Dask tasks for each chunk of the final array, following the example here. This shaves a bit of time off when using the distributed scheduler (36 to 32 seconds), but the "processes" scheduler is still faster and is almost unaffected by this change (from a consistent 20.5 to 19.7 seconds). I'm using the same test case as in #455 .

This PR currently needs more thought on handling non-default chunksizes (the argument to stack_loader_array, which it looks like can be set in the ASDF file?). I guess I'm not sure what the intended use case is for that. In main it looks like this option can only sub-divide files into smaller chunks. Are there DKIST files out there that are too big to be a single chunk? Or is the idea to subdivide files when only a portion of each file has to be loaded? As this PR is now, each file gets fully loaded, and then split up, which might be less efficient. But before this PR, I'm not sure if that was also the case, or if hdu.data[slice] loads just the sliced portion of the file. If the latter, changing this PR to make use of that would, I think, require manually creating tasks for each chunk of each file, which seems complicated.

In #455, there was discussion of rechunking to larger chunk sizes. At least for that test case, using .rechunk((1, 100, -1, -1)) actually slows down the total computation (I'm seeing it go from 36 to 45 seconds), I assume from the overhead of gathering and combining arrays. This PR's design makes it easy to try larger chunk sizes more efficiently, as one "loading" task can load multiple files and concatenate them immediately, rather than adding a separate layer to the task graph. This is implemented in this PR, in the Docstring commit and before. A "batch size" can be set, and each Dask task loads that many files and concatenates them. (The batch size has to be adjusted in the code---I didn't know how to expose it as an option for load_dataset.) I found that with a batch size of 100 (producing chunks of a few hundred MB, Dask's recommended size), I was actually getting a very slightly longer runtime (going from 24 to 27 seconds), but it's less extreme than using .rechunk(). My guess is that the np.concatenate call copying all the data arrays offsets any efficiency gains from the larger chunks. (Maybe if I were doing a lot more with the data it would be worth it.) I didn't see a way to have Astropy FITS load a file directly into a view of a larger output array, unfortunately.

Here's the task stream (top is before, bottom is this PR):

And here's a portion of the task graph for each (left is before, right is this PR). This is from dataset.data.visualize(), and I made an asdf file with only a handful of FITS files so the images wouldn't be unrenderably large.

Right now my thinking is "just use the processes scheduler", where the gains from this PR are small enough that I'm not sure it's worth more work to handle the chunksize argument better.