struggles with dread and long files

[ryan-budde]

@cwindolf along with everyone else

As discussed in the NPX discord I've settled on some dredge parameters that I'm hopeful will give me good results, doing extensive testing on data snippets ranging 3-10 minutes.

Now I am attempting a run on my whole experiment file which is 12868 seconds (3.5 hours).

I am getting a lot of OOM errors in dredge and I'm working through them and would appreciate some help. A small one was in the laplacian function, at the step lap = np.zeros((n, n)) and I OOM'd needing only a MB, so I added a garbage collection and that worked and has never erred again. To me that's a normal OOM since scipy / python is bad with memory leaks and when processing these large loops garbage collection is required, and you never know where it will come up until you run long files. This is not an error that concerns me for myself or for other users.

My main issue, and what I anticipate will cause issues for many other users, are the large arrays Cs Ds Us and Ss, which are not memory leak issues but rather huge arrays that cannot be reasonably loaded into memory.

Ds and Cs are defined with Ds = np.zeros((B, T0, T1), dtype=np.float32) where, for me, it looks like both T0 and T1 are 12868 seconds, and B is my total number of spatial windows, which is 148. For float 32 this works out to 98 GB (each array) RAM or on my hard drive.

Later Ss is a modified copy of Cs and Us is a modified copy of Ss. So for my file it is 98*4 nearly 400 GB of required RAM. But I believe it is possible to make no more than 3 active copies at a time, since I believe Cs is deleted after Ss is calculated, and Ss is deleted after Us is calculated. I don't think this is implemented yet and all four copies are in memory at some point.

In the final thomas_solve step both Us and Ds are used and are upgraded to float64. This is requiring something like 196 GB of memory (per array), which is not reasonable for a single machine, even professional grade. And I've definitely run longer experiments before, 5 to 6 hours, which would require something like 1.1 TB of memory (just for these two arrays).

My current solution that I am hacking into the code is to use memmaps and write these large arrays to the disk. I am still getting errors and I do not have a functional version yet, but this should work "in theory" and simply requires a large, fast hard drive of sufficient space.

However my intuition is that there is a more optimal solution for long experiments where these arrays do not fit into memory. I think the current code is likely speed optimal, but speed optimal is not useful if it completely fails on the dataset I need to use. And I think that in general most NPX users are using much longer files that will reasonably OOM with the current implementation. Assuming a user has only 64 GB of RAM, then in the final thomas_solve step (two arrays Us and Ds, 64 bits, of size 148xTxT) they can only run a file of length 85 minutes (by my math, assuming no other memory usage beyond these variables).

Are there any ways that the speed-optimal code can be modified to be better than a memmap and still run for long files? In particular I am interested in why the entire BxTxT array needs to be in memory at one time. Especially since I have modified my time horizon to only be about 4 seconds, I would imagine that the code could be restructured to calculate motion along a sliding time window, something like BxTHxTH instead? Or perhaps the steps to go from Cs to Ss to Us could be calculated "on the fly" for each B segment of data? Of course these implementations are much slower as they would involve repeat calculations that are then discarded, but for me (and I assume many others) the file length limit is a deal-breaker, and so speed on shorter files is irrelevant if longer files completely fail.

[yger]

I guess indeed, we should improve such a method. As already discussed, and I think @cwindolf approved, most of the internal data structures in dredge could be float32 only, to reduce memory footprint. This would clearly not solve all your problems, but this is a start.

Regarding the time horizon, I agree, but I need to read the code more in depth to see how it could be implemented, also maybe with some sparse data structures

[cwindolf]

Hi @ryan-budde, thanks for your detailed investigation! I have a couple of ideas.

• Windows: First, I want to double check that you really need to use 148 spatial windows. Is there fine spatial structure in the motion that need to be modeled at such a high spatial resolution? (Assuming NP1, 3840/148 = 25 um resolution). The reason I ask is that, for reference, I have used at most 10 windows in NP1 data, although I have needed narrower windows for NP Ultra or for 4-shank NP2 data in certain settings. Still, I haven't come across an example needing 25um resolution and my expectation would be that you can probably raise that. If you're able to raise the win_scale_um (maybe without therefore also raising max_disp_um), you may be able to cut down the window count drastically, I'd expect by a factor of 5 or 10?
• float32: It would probably be possible as @yger suggests to store everything as float32, but I'd just want to make sure that they're back to double just before any calls to solve(). But only a couple of TxT quantities would need to be double at any moment.
• Sliding time window: If the above are not enough to get the method working for you, you could test out a sliding time window strategy based on the online method we use for LFP. The idea would be to put your spike raster into a form that can be input into https://github.com/evarol/dredge/blob/main/python/dredge/dredge_lfp.py#L10 . If you're interested in this idea, I can be more specific about how it would work, and we could discuss over Discord or a call sometime. It could turn into something that we'd submit as a PR to SpikeInterface as a new dredge_online motion estimation strategy.
  • This is something I've been meaning to implement for a while, in particular for chronic recording use cases. It has the nice property that adding more time to your recording would not cause DREDge to revise its motion estimate for earlier times in the recording.

Considerations about some other points that were raised:

• I was considering whether it would be possible to avoid storing B TxT matrices (whether in memory or on disk) at the same time. Could you just stream along the B dimension? Sadly no, because we need to have all of the "gamma_hat" matrices ready for the back substitution.
• Re storing just the time horizon bands, I think it's a nice idea. At some point I investigated using https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.solve_banded.html but found that regular solve() has better stability. Still, as with the float32/double point raised above, you could store just the banded form and convert to dense as needed; sadly, though, I'm not at the moment sure if there is any guarantee that the gamma_hat matrices would also be sparse/banded, so it may be necessary to store at least one BxTxT thing.