Cache some charge increment calculations

[mattwthompson]

Description

For some systems, most of the time writing GROMACS files was spent adding up a bunch of charges. This is a shortcut to speed that up. Some rough checks using some of the slower toolkit examples:

$ python -m pytest --nbval-lax --dist loadscope -n logical --durations=20 ../openff-toolkit/examples/using_smirnoff_in_amber_or_gromacs ../openff-toolkit/examples/using_smirnoff_with_amber_protein_forcefield -p no:randomly > head.log

With e6f9f88:

============================= test session starts ==============================
platform darwin -- Python 3.11.11, pytest-8.3.4, pluggy-1.5.0
rootdir: /Users/mattthompson/software/openff-toolkit
configfile: pyproject.toml
plugins: cov-6.0.0, nbval-0.11.0, anyio-4.6.2.post1, xdist-3.6.1
created: 8/8 workers
8 workers [25 items]

.........................                                                [100%]
============================= slowest 20 durations =============================
119.34s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 9
82.66s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 3
49.39s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 8
49.14s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 10
48.09s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 5
29.25s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 6
8.10s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
7.50s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 4
6.49s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 1
3.43s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 3
2.03s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 2
0.85s setup    examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 0
0.85s setup    examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
0.69s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 6
0.53s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 1
0.45s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 2
0.37s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 8
0.30s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 12
0.27s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 11
0.16s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 9
======================== 25 passed in 396.25s (0:06:36) ========================

Compared to 0.4.0:

============================= test session starts ==============================
platform darwin -- Python 3.11.11, pytest-8.3.4, pluggy-1.5.0
rootdir: /Users/mattthompson/software/openff-toolkit
configfile: pyproject.toml
plugins: cov-6.0.0, nbval-0.11.0, anyio-4.6.2.post1, xdist-3.6.1
created: 8/8 workers
8 workers [25 items]

.........................                                                [100%]
============================= slowest 20 durations =============================
516.67s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 9
257.41s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 10
245.22s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 8
67.91s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 3
36.93s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 5
20.93s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 6
6.45s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 3
6.31s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 4
1.28s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
1.26s setup    examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
1.26s setup    examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 0
0.91s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 1
0.88s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 6
0.79s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 2
0.59s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 1
0.52s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 12
0.27s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 2
0.27s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 11
0.22s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 8
0.12s teardown examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 13
======================= 25 passed in 1155.80s (0:19:15) ========================

Checklist

• Benchmark .to_gromacs on several systems
• Add tests
• Lint
• Update docstrings

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 93.38%. Comparing base (94c1ad1) to head (fe8f88f).
Report is 14 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1126      +/-   ##
==========================================
+ Coverage   93.36%   93.38%   +0.01%     
==========================================
  Files          69       69              
  Lines        5997     5997              
==========================================
+ Hits         5599     5600       +1     
+ Misses        398      397       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mattwthompson]

Some nice improvements for a couple of systems (GB3 down to ~2 seconds from 12-18)

This script should be gobbled up into a solution for #1020

import warnings

warnings.filterwarnings('ignore')

from openff.toolkit import ForceField, Molecule, Topology, Quantity
from IPython import get_ipython

gb3 = Topology.from_pdb("../proteinbenchmark/proteinbenchmark/data/pdbs/gb3-1P7E.pdb")
gb3.box_vectors = Quantity([10, 10, 10], "nanometer")

protein = ForceField('Protein-Specific-0.0.3.offxml').create_interchange(gb3)

caffeine = Molecule.from_smiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
caffeine.generate_conformers(n_conformers=1)
caffeine = caffeine.to_topology()
caffeine.box_vectors = Quantity([10, 10, 10], "nanometer")

ligand = ForceField("openff-2.2.1.offxml").create_interchange(caffeine)


print("Exporting ligand system:")
get_ipython().run_line_magic(
    "timeit",
    "ligand.to_gromacs('tmp')",
)

print("Exporting protein system:")
get_ipython().run_line_magic(
    "timeit",
    "protein.to_gromacs('tmp')",
)

$ ipython gmx_perf.py && git checkout upstream/main && ipython gmx_perf.py && git checkout v0.4.0 && ipython gmx_perf.py
Exporting ligand system:
9.48 ms ± 496 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Exporting protein system:
2.15 s ± 72.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
HEAD is now at 3e65e567 Merge pull request #1072 from openforcefield/dependabot/github_actions/mamba-org/setup-micromamba-2
Exporting ligand system:
16.1 ms ± 277 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Exporting protein system:
11.5 s ± 269 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Previous HEAD position was 3e65e567 Merge pull request #1072 from openforcefield/dependabot/github_actions/mamba-org/setup-micromamba-2
HEAD is now at 7f93571f Merge pull request #1090 from openforcefield/fix-missed-decorator

Support for writing out LAMMPS data files will be removed
in mbuild 1.0.
See GMSO (https://github.com/mosdef-hub/gmso/tree/main/gmso/formats/lammpsdata) for
continued support for LAMMPS.


Exporting ligand system:
17.8 ms ± 3.29 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
Exporting protein system:
11.4 s ± 808 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[mattwthompson]

This conflicts a little with #1072, I'm not sure which is preferred

[mattwthompson]

============================= test session starts ==============================
platform darwin -- Python 3.11.11, pytest-8.3.4, pluggy-1.5.0
rootdir: /Users/mattthompson/software/openff-toolkit
configfile: pyproject.toml
plugins: cov-6.0.0, nbval-0.11.0, anyio-4.6.2.post1, xdist-3.6.1
created: 8/8 workers
8 workers [19 items]

...................                                                      [100%]
============================= slowest 20 durations =============================
25.20s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 2
4.22s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 4
2.28s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 3
1.19s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
1.12s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 5
0.89s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 3
0.84s setup    examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
0.84s setup    examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 0
0.77s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 1
0.57s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 6
0.40s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 0
0.37s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 1
0.30s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 12
0.22s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 11
0.21s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 5
0.19s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 8
0.16s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 2
0.14s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 10
0.11s teardown examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 12
0.11s teardown examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 5
============================= 19 passed in 34.71s ==============================

[j-wags]

This conflicts a little with #1072, I'm not sure which is preferred

IIRC I was assigned this PR to answer this question, but the wrong thing is linked here and I'm not easily able to find what was supposed to be linked. Could you check this?

[mattwthompson]

Yep that's the wrong number, I was fishing for #1115 but they don't intersect as obviously as I previously thought

[mattwthompson]

Using the toolkit's examples as a benchmark was a shifting target because of a major change (openforcefield/openff-toolkit#1981).

Re-running the same command the latest commit (6fc864c):

============================= test session starts ==============================
platform darwin -- Python 3.11.11, pytest-8.3.4, pluggy-1.5.0
rootdir: /Users/mattthompson/software/openff-toolkit
configfile: pyproject.toml
plugins: cov-6.0.0, nbval-0.11.0, anyio-4.6.2.post1, xdist-3.6.1
created: 8/8 workers
8 workers [21 items]

.....................                                                    [100%]
============================== slowest durations ===============================
22.57s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 3
2.09s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 3
1.82s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 6
0.91s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
0.79s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 4
0.77s setup    examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 0
0.77s setup    examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
0.54s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 1
0.49s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 0
0.48s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 6
0.24s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 12
0.22s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 1
0.21s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 11
0.20s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 8
0.18s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 2
0.11s teardown examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 6
0.11s teardown examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 13
0.08s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 10
0.04s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 5
0.04s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 4
0.03s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 9
0.02s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 7

(41 durations < 0.005s hidden.  Use -vv to show these durations.)
============================= 21 passed in 27.90s ==============================

compared on the `main` branch (b623003)

============================= test session starts ==============================
platform darwin -- Python 3.11.11, pytest-8.3.4, pluggy-1.5.0
rootdir: /Users/mattthompson/software/openff-toolkit
configfile: pyproject.toml
plugins: cov-6.0.0, nbval-0.11.0, anyio-4.6.2.post1, xdist-3.6.1
created: 8/8 workers
8 workers [21 items]

.....................                                                    [100%]
============================== slowest durations ===============================
25.00s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 3
3.34s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 3
2.21s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 6
1.37s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
1.13s setup    examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 0
1.13s setup    examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 0
1.02s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 1
0.83s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 4
0.66s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 6
0.35s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 0
0.31s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 8
0.27s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 1
0.22s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 12
0.21s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 11
0.19s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 2
0.17s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 5
0.12s teardown examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 13
0.11s teardown examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 6
0.11s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 10
0.06s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 4
0.03s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 9
0.02s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 7
0.01s call     examples/using_smirnoff_in_amber_or_gromacs/export_with_interchange.ipynb::Cell 13
0.01s call     examples/using_smirnoff_with_amber_protein_forcefield/BRD4_inhibitor_benchmark.ipynb::Cell 2

(39 durations < 0.005s hidden.  Use -vv to show these durations.)
============================= 21 passed in 31.67s ==============================

The speedup is relatively minor but I think still worthwhile. Open to feedback

[mattwthompson]

As long as the macOS jobs continue to come back green, I think this is good to go

[j-wags]

Looks great!