Repository navigation
depletion: fix performance of chain matrix construction - #3567
Conversation
The previous construction was repeatedly invoking _validate_indices for the scipy sparse matrix which does tons of expensive checks for every square bracket indexing operation. This was expensive and was causing my full core depletion runs to take 10 hours with 90% of the time spent in these validate indices checks. This skips that and does a batch/all-at-once construction of the sparse matrix after accumulating all the indices and values.
|
Super nice. Would you want to try pre-building the sparsity pattern then running this? I bet that would be even faster. This is currently probably still allocating on-the-fly as you loop over the different terms. Eh, considering the extra complexity, maybe that's not really worth it. |
paulromano
left a comment
There was a problem hiding this comment.
@rwcarlsen Thanks a lot for this improvement. It's kind of mind-blowing to me that indexing on dok_matrix would be that slow, but in any case I like the new approach. The one part that I'm slightly wary of is that setval(i, i, ...) may be called twice; this happens when a nuclide is radioactive but also has reactions that can transmute it. Evidently, passing (vals, (rows, cols)) to csc_matrix where there are row,col duplicates still works --- when a (row, col) pair appears twice, it will add on to a value that is already there. For example,
>>> import scipy.sparse as sp
>>> sp.csc_matrix(([0.1, 0.1], ([1, 1], [1, 1])), shape=(3, 3)).todense()
matrix([[0. , 0. , 0. ],
[0. , 0.2, 0. ],
[0. , 0. , 0. ]])The part that worries me is that I don't see anything in the documentation that provides a guarantee of this behavior, so it's conceivable that scipy could change that behavior in the future (those who haven't heard of Hyrum's law should look it up). Perhaps a better option here would be to build a dictionary mapping {(row, col): val} (with a check for existing keys) and then iterate to build the necessary lists at the end before constructing csc_matrix.
@GuySten Thanks a lot for the quick review here. However, please be aware that per our contribution guidelines, PRs should not be merged for at least 36 hours after they were originally submitted to allow time for people to chime in. Partly my fault for not making this clear when I gave you merge rights 😄 In the future, please be mindful of that, and particularly if there are PRs related to performance, I'd appreciate a ping before you hit the big green button.
|
@paulromano, I will keep that in mind. |
Description
The previous construction of the chain sparse matrix was repeatedly invoking square bracket indexing which calls into the expensive
_validate_indicesfunc. This was causing my full core depletion runs to take 10s of hours with ~90% of the time spent in these validate indices checks. This avoids that and does a batch/all-at-once construction of the sparse matrix after accumulating all the indices and values.Sorry for the quick drive-by. I haven't really looked over your contributor guidelines or done any due diligence beyond ensuring this makes my depletion runs ~10x faster. Figured it might be consequential enough that you'd want to at least see the issue.
Checklist