More efficient cohorts.

[dcherian]

Closes #140

We apply the cohort "split" step after the blockwise reduction, then use
the tree reduction on each cohort.

We also use the .blocks accessor to index out blocks. This is still a
bit inefficient since we split by indexing out regular arrays, so we
could index out blocks that don't contain any cohort members. However,
because we are splitting after the blockwise reduction, the amount of
work duplication can be a lot less than splitting the bare array.

One side-effect is that "split-reduce" is now a synonym for "cohorts".
The reason is that find_group_cohorts returns a dict mapping blocks to
cohorts. We could invert that behaviour but I don't see any benefit to
trying to figure that out.

[dcherian]

For this dataset:

I get (sort=False)

|                       | unoptimized | optimized |
|-----------------------+-------------+-----------|
| main cohorts          |        9563 |      1307 |
| this PR cohorts       |        6866 |      1351 |
| this PR after c48041c |        6425 |      1221 |
| map-reduce            |        5638 |      1018 |

Looking at graphs for the first two blocks:

• This branch must show lower memory usage but I need a benchmark for this
• This branch has a lot of blocks tasks because I extract blocks in a loop for advanced indexing. +I think that could be optimized a bit to reduce how much I loop (but not too hopeful)+
• I could optimize to do a blockwise thing instead of tree-reduce where that makes sense but seems minor.
• Note that by a simple num-tasks metric, map-reduce looks great but it involves a lot of unnecessary communication. We are trading off larger number of tasks for fewer transfers and communication load.

This branch graph

main branch graph

[dcherian]

For this array and reducing over the counties in the previous image:

|                  | map-reduce | cohorts   |
|------------------+------------+-----------|
| number of tasks  | 1856       | 6420      |
| compute time     | 2372.83 s  | 2075.00 s |
| deserialize time | 4.68 s     | 2.87 s    |
| transfer time    | 15.18 s    | 5.17 s    |
| total transfer   | 469.56 MB  | 164.44 MB |

So! transferring less data, and spending less time on those transfers, exactly as hoped. This was on my laptop so in theory, this might make a difference on cloud/HPC systems.

This is with reindex=True. I did find map-reduce with reindex=False to be faster, but with similar network transfers.

The compute time is massively dominated by the blockwise reduction with numpy-groupies, so optimizing that would be very impactful!