Use cache to speed up applying forcefield to connections

[chrisjonesBSU]

PR Summary:

I've been doing some digging into different ways to speed up GMSO, especially related to applying a force field.

This PR adds a signature cache to _apply_connection_parameters that lets us skip get_connection_identifiers for a connection that has already been found (if it has the same sig-cache). This is especially useful for molecules that have a bunch of repeated bonds, angles and dihedrals (like polymers).

This script below with 500 alkane chains currently takes 160 seconds to apply the forcefield, but only 7 seconds with these changes!

import mbuild as mb
import gmso
from gmso.parameterization.parameterize import apply
from gmso.core.forcefield import ForceField
import numpy as np

alkane = mb.load("CCCCCCCCCCCCCC", smiles=True)
alkane.name = "alkane"
comp = mb.fill_box(
    compound=[alkane], n_compounds=[500], density=0.2
)
top = comp.to_gmso()
oplsaa = ForceField("oplsaa")

start = time.time()
typed_top = apply(
    top=top,
    forcefields={"alkane": oplsaa},
    identify_connections=True,
    speedup_by_moltag=True,
)
finish = time.time()
total = finish - start
print(f"Finished in {np.round(total, 3)} seconds")

There are a couple other speed up additions:

1. Similar to some recent PRs where we build up a dict ahead of time for lookup rather than for loops, molecule_utils.py gets a method that does this for connections. This is used when iterate through multiple force fields passed to apply(), so the speed up only occurs in that scenario. This resulted in ~2x speed up for small systems. In larger systems, other areas dominated the time required.

2. connections_identifier now acts as a generator rather than building and returning a list.

PR Checklist

---

• Includes appropriate unit test(s)
• Appropriate docstring(s) are added/updated
• Code is (approximately) PEP8 compliant
• Issue(s) raised/addressed?

[codecov]

Codecov Report

❌ Patch coverage is 88.88889% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.33%. Comparing base (1b8bb2c) to head (5099f18).
⚠️ Report is 60 commits behind head on main.

Files with missing lines	Patch %	Lines
gmso/parameterization/topology_parameterizer.py	84.21%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #978      +/-   ##
==========================================
- Coverage   93.36%   93.33%   -0.03%     
==========================================
  Files          67       67              
  Lines        8066     8105      +39     
==========================================
+ Hits         7531     7565      +34     
- Misses        535      540       +5     

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[CalCraven]

I think we can merge this with some minor changes, but I do think we need to relook at the molecule speedups in the topolgoy parameterizer to see if something is failing there, or potentially to merge these methods so they work together better. Some of this molecule identification should happen in certain gmso systems using the networkx graph.

[CalCraven]

I think this works as is because site.atom_type.name is always required, but site.atom_type.atom_class is not always there. However, a lot of times the atom_classes are a more general representation of the connection. Wonder if we wanted even more speedup if you could check for the classes here instead of the types?

[chrisjonesBSU]

Ok, this should be added now.

[CalCraven]

As far as I'm aware, we don't have any test examples for systems with bonds that cross a molecule. Or any tests for using "group" as the connection identifier.

I think we could use some sort of molecule_validation method, which essentially is doing this check. Can we rename this function to validate_and_bucket_molecules?
Then we can raise an error right here if something doesn't pass that check, instead of just silently skipping it.

What do you think?

[chrisjonesBSU]

Good catch. I think the actual validation is done before hand by the below method, which is called in topology_parameterizer.py a couple lines before build_molecule_connection_index

def assert_no_boundary_bonds(top):
    """Assert that all bonds in the topology belongs to only one molecule."""
    assertion_msg = "Site {} is in the molecule {}, but its bonded partner {} is in the molecule {}."
    for bond in top.bonds:
        site1, site2 = bond.connection_members
        assert site1.molecule == site2.molecule, assertion_msg.format(
            site1.name, site1.molecule, site2.name, site2.molecule
        )

So, I think this if statement isn't needed. We could just remove it.

Ultimately, I think this raises questions about rigidity of hierarchy in gmso. Doesn't a method like assert_no_boundary_bonds prevent using FFs at a residue level, where one part/residue of a molecule uses one FF, and another uses a different one? I don't think this is an issue we need to solve in this PR. For now, I'll just remove the if statement.

[chrisjonesBSU]

I think we can merge this with some minor changes, but I do think we need to relook at the molecule speedups in the topolgoy parameterizer to see if something is failing there, or potentially to merge these methods so they work together better. Some of this molecule identification should happen in certain gmso systems using the networkx graph.

There are some operations that scale poorly with the topology size even when using speed up by mol tags (like _apply_connection_parameters). I think speedup_by_moltag mostly only applies to atom typing. Also, I think we should be selective about when and where we rely on graph-matching for speed-up purposes. I don't think the networkx isomorphism checks are necessarily that fast, especially as the size of the graph increases. Their usefulness mostly comes from allowing flexible/arbitrary graph comparisons (like atom typing), but not so much when the graph structure is strictly defined (like comparing connections) where regular-old algorithms might be better, like we saw in #972 .

I think one approach to this might be something I was working on in #961. Things like applying a forcefield, identifying connections, calling charge calculation methods (the goal of that PR), writing out hoomd forcefield objects really only need a minimum representation of the topology (i.e., a single molecule, or the set of unique molecules in a non-homogenous system). But we can discuss that more in a separate issue.