proteneer / timemachine

Differentiate all the things!
Other
138 stars 17 forks source link

Add Potential::execute_batch_sparse #1308

Closed mcwitt closed 2 months ago

mcwitt commented 3 months ago

Adds a "sparse" variant of the existing Potential::execute_batch function that can be used to evaluate a subset of elements of the dense (coords_batch_size, params_batch_size) matrix. Compared to execute_batch, execute_batch_sparse additionally accepts two arrays (of length batch_size) specifying the index of the coordinates and parameters, respectively, to use for each evaluation.

The present motivation is to enable more efficient computation of the $u_{kl}$ matrix used in HREX, since most of the energies in the dense matrix end up unused (we typically only need energies in the neighborhood of each replica).

mcwitt commented 2 months ago

Converting back to draft as the only current application of this in #1309 does not (initially) appear to have a significant performance impact; need to investigate further.

mcwitt commented 2 months ago

Reopening for review as more careful benchmarks do show a significant improvement in #1309.