Sub-batches with equal flux in photon pooling

[welucas2]

Currently, the sub-batching in photon pooling is performed object-by-object. So, for example, if there are 100 objects in the batch, and 20 sub-batches, then the first five objects will go in the first sub-batch, the next five in the second, and so on. This means that if the objects vary a lot in brightness (and they do) then the flux per sub-batch will also vary a lot. In a typical run and setup provided by Jim, I found that most sub-batches contained a few thousand photons, and a handful contained several tens of millions. This leads to regular memory spikes when those sub-batches are processed in each new photon pool.

This PR smooths those memory spikes out by creating sub-batches with roughly equal fluxes. This is an implementation of the bin packing problem with fragmentation: our goal is to split the total flux in the pool across the sub-batches, and to achieve this we have to be able to fragment those extremely bright objects across however many sub-batches.

At the same time, each object fragmentation means another object lookup, and we had determined previously that with photon pooling these can mount up and become quite costly. For this reason, I'm allowing the flux in each sub-batch to vary slightly, letting objects to fill up over the limit to 105% of the expected per sub-batch flux if it prevents a fragmentation.

This is ready to go, but needs #511 to be merged in first and then this should be rebased onto it - so for now I'm leaving this as a draft.

[welucas2]

This is ready for review, though CI failed earlier on today during Conda setup: CondaHTTPError: HTTP 000 CONNECTION FAILED for url <https://conda.anaconda.org/conda-forge/linux-64/repodata.json>. Could someone with permission please re-run CI to see if this gets past it?

[rmjarvis]

I think I'm confused about some of this.

[rmjarvis]

I looked at this again, and I realized that the algorithm for splitting up the photons didn't do a great job when there was a large (and smooth) dynamic range in the flux values. I added a test that uses a power law distribution for the fluxes, so there are a few bright things that need to be split and a tail all the way down to f=1 objects, which are much more numerous. This broke the assert max < 1.1 * min test that was in there. This distribution tended to have one subbatch with much less flux in it. Like 0.3 of the average.

I changed the algorithm a fair bit. It now passes all the old tests and it ends up essentially completely flat in nphot per subbatch for all the test distributions, including the power law. The basic idea is to always switch to the subbatch with the most space whenever you get an object that doesn't fit in the current one. Then if it still doesn't fit, split it. Since we deal with objects in decreasing order of flux, this means that it only splits objects that really don't fit anywhere. I also lowered the amount of spill over allowed to 1% of the mean, rather than 5%, since it's not nearly as important for this algorithm to prevent spurious splitting.

[jchiang87]

I ran a couple of tests using a grid of bright stars, configured so that photon shooting was always used, and monitored the memory using in each case. Using current main, I see memory spikes up to 18GB:

With this branch, memory usage is smoothed out considerably, the "baseline" memory is lower overall, and the run time is actually somewhat shorter, 7.0 ks vs 7.4 ks using main:

I had a quick look at the code as well, and everything looks good to me, so I'd say we should go ahead and merge this branch.