proposal: new distributed histogram type

[lgray]

I've attached a notebook showing a proof-of-concept for a cluster-distributed histogram using dask-array that relies on the sparseness of per-chunk fills in dask to avoid materializing the histogram on any particular worker.
This implementation makes a number of shortcuts and fusions of operations that building this out of the shuffle primitives in dask lays out a bit more verbosely.

Right now this only works for unweighted histogram and no categories but it's not too much of a stretch to extend it. The main trouble will be keeping the number of additionally created tasks down per category.

I have scaled to decent performance for O(billion) bin histograms (checked up to 4D) and it all fits on my laptop and fills in a reasonable amount of time compared to da.histogramdd and with significantly less memory usage than the present implementation of dask histogram's boost backend. The histogram can then be reassembled on the client machine and quickly turned into boost-histogram or anything else suitable.

Here's the notebook with the proof of concept and validation against da.histogramdd:
distributed-histogram.ipynb.zip

Here are settings for scaling to very large number of bins successfully:
NOTE: The da.histogramdd check will perform extremely sluggishly for the following settings
125M bins:

ndims = 3
nbins = 500
hist_partitions = 50
n_hist_slices = 25

1B bins:

ndims = 4
nbins = 178
hist_partitions = 100
n_hist_slices = 10

These settings will run pleasantly on a modern laptop. You're more or less limited by system memory for the output histogram.

@nsmith- @holzman @mapsacosta @Nanoemc @douglasdavis @martindurant
Thoughts on how to improve or integrate with existing dask-histogram greatly appreciated.

Jotting down my own thoughts concerning moving forward:
For weighted filling we need to keep track of a tuple of (filled_bins, sum_weights, sum_weights_squared). That's not exactly congruent with the dask-array setup I'm using right now, but it is still congruent with blockwise, piping it all together sounds a little nasty.

For categories we'll need to be careful about when we actually make more tasks since it's easy to generate an enormous task overhead with a category for each systematic variation. However, since we know when multiple .fill calls take the same coordinate collections and only a different weight collection and category label, we can figure out how to fuse those when building the task graph, and then do a vectorized fill over weights for things like that and not actually bloat the size of the graph.

[martindurant]

For weighted filling we need to keep track of a tuple of (filled_bins, sum_weights, sum_weights_squared)

Can it be done by using a complex dtype (i.e., record array) ?

[lgray]

I don't think you even need record arrays (at some slight debugging inconvenience).

I was considering this morning that you can manage it if you go to unknown lengths and then store the sparse histogram as the (N, 3) array [bin_idxs, sum_w, sum_w2]. Since the non-summed bincounts are intermediates and not seen by the user and you know the output shape and chunking exactly some random unknown chunk size array isn't much of an issue.

Though this also begs the question of dask-array being the right thing to implement this in! It certainly performs the task.

Dask-awkward is an interesting possibility but not sure yet if necessary. It might be what's required for managing the categories in a human-comprehendible way?

[lgray]

Ah, yeah I guess it needs a RecordArray or something similar since the indices will usually have a different dtype from the data.

It seems like just good old numpy structured arrays will work?

[lgray]

OK - I've got it working with weights to no unexpected outcome. This implementation uses a structured array dtype as an intermediate dataformat, and doesn't use a sparse matrix class anymore. Instead we use a bare bones indices-and-values representation. The implementation tracks sumw and sumw2 throughout the calculation and returns the filled histogram as a (nbins, ..., 2) shaped output array. It could be made more numerically reproducible if I sorted ascending by the weights before filling.

Still scales to a very large number of bins, the total memory use goes up by 2x to track the (effects of) weights.

distributed-histogram-weights.ipynb.zip

Next steps are to wrap this up into a generating function and class and start getting it closer to actual other histogram implementations. After that, I'll start working on categories which need some care.

[martindurant]

Feel free to add the ipynb file as an example in the PR, which will eventually be good for docs and allow github to directly render it.