Luthaf / rascaline

Computing representations for atomistic machine learning
https://luthaf.fr/rascaline/
BSD 3-Clause "New" or "Revised" License
44 stars 13 forks source link

Reduce the number of equistore access in spherical expansion #146

Closed Luthaf closed 1 year ago

Luthaf commented 1 year ago

Previously, we where looking up block/sample/property for each pair and directly storing pair contributions into equistore. This lookup and accumulation was taking 4x longer than the calculation of pair contribution to the spherical expansion.

The new code accumulates data for all pairs first into temporary arrays, and later moves the data to equistore, making the spherical expansion up to 60% faster without gradients and 15% faster with gradients; while increasing the max RSS requirements by around 35%.

This was discovered while porting rascaline to use https://github.com/lab-cosmo/equistore/pull/85. I'll check that it does improve the results there as well before merging this PR!

benchmarks results

``` group fewer-equistore-in-spx master ----- ---------------------- ------ SOAP PS(per atom) with gradients/Bulk Silicon/n_max = 15, l_max = 14 1.00 720.2±4.79µs 1.11 800.6±8.17µs SOAP PS(per atom) with gradients/Bulk Silicon/n_max = 2, l_max = 1 1.00 15.2±0.23µs 1.00 15.2±0.75µs SOAP PS(per atom) with gradients/Bulk Silicon/n_max = 8, l_max = 7 1.00 104.9±0.46µs 1.04 109.5±0.66µs SOAP PS(per atom) with gradients/Molecular crystals/n_max = 15, l_max = 14 1.00 2.2±0.06ms 1.05 2.3±0.03ms SOAP PS(per atom) with gradients/Molecular crystals/n_max = 2, l_max = 1 1.00 39.7±0.53µs 1.01 40.3±0.36µs SOAP PS(per atom) with gradients/Molecular crystals/n_max = 8, l_max = 7 1.00 277.8±4.47µs 1.02 284.7±6.89µs SOAP PS(per atom)/Bulk Silicon/n_max = 15, l_max = 14 1.00 44.2±0.64µs 1.42 62.6±1.39µs SOAP PS(per atom)/Bulk Silicon/n_max = 2, l_max = 1 1.00 1951.1±62.96ns 1.09 2.1±0.02µs SOAP PS(per atom)/Bulk Silicon/n_max = 8, l_max = 7 1.00 6.8±0.09µs 1.37 9.2±0.06µs SOAP PS(per atom)/Molecular crystals/n_max = 15, l_max = 14 1.00 134.4±4.91µs 1.24 166.9±4.07µs SOAP PS(per atom)/Molecular crystals/n_max = 2, l_max = 1 1.00 9.5±0.13µs 1.04 9.9±0.11µs SOAP PS(per atom)/Molecular crystals/n_max = 8, l_max = 7 1.00 28.3±0.21µs 1.18 33.5±0.37µs SOAP SPX(per atom) with gradients/Bulk Silicon/n_max = 15, l_max = 14 1.00 539.9±10.25µs 1.12 605.9±5.31µs SOAP SPX(per atom) with gradients/Bulk Silicon/n_max = 2, l_max = 1 1.00 11.1±0.63µs 1.03 11.4±0.39µs SOAP SPX(per atom) with gradients/Bulk Silicon/n_max = 8, l_max = 7 1.00 77.8±0.38µs 1.07 83.4±0.40µs SOAP SPX(per atom) with gradients/Molecular crystals/n_max = 15, l_max = 14 1.00 816.8±28.10µs 1.15 935.6±6.55µs SOAP SPX(per atom) with gradients/Molecular crystals/n_max = 2, l_max = 1 1.00 14.7±0.21µs 1.09 16.0±0.25µs SOAP SPX(per atom) with gradients/Molecular crystals/n_max = 8, l_max = 7 1.00 107.5±1.81µs 1.08 116.2±2.82µs SOAP SPX(per atom)/Bulk Silicon/n_max = 15, l_max = 14 1.00 32.7±0.50µs 1.68 55.0±1.62µs SOAP SPX(per atom)/Bulk Silicon/n_max = 2, l_max = 1 1.00 1448.2±33.84ns 1.10 1589.8±20.90ns SOAP SPX(per atom)/Bulk Silicon/n_max = 8, l_max = 7 1.00 7.2±0.04µs 1.19 8.6±0.16µs SOAP SPX(per atom)/Molecular crystals/n_max = 15, l_max = 14 1.00 51.0±1.60µs 1.63 83.0±5.17µs SOAP SPX(per atom)/Molecular crystals/n_max = 2, l_max = 1 1.00 4.2±0.04µs 1.09 4.6±0.06µs SOAP SPX(per atom)/Molecular crystals/n_max = 8, l_max = 7 1.00 11.1±0.12µs 1.46 16.2±0.30µs ```
github-actions[bot] commented 1 year ago

The documentation for this PR is (or will soon be) available on readthedocs: https://rascaline--146.org.readthedocs.build/en/146/

github-actions[bot] commented 1 year ago

Here is a pre-built version of the code in this pull request: wheels.zip, you can install it locally by unzipping wheels.zip and using pip to install the file matching your system