Improve speed of OFF molecule/network molecule comparisons and isomorphisms

[j-wags]

Description

System is a box of 1000 hexadecanes.

from openforcefield.typing.engines.smirnoff import ForceField
from openforcefield.topology import Molecule, Topology
from simtk.openmm import app
import time

hexadecane = Molecule.from_smiles('CCCCCCCCCCCCCCCC')
pdbfile = app.PDBFile('output.pdb')
start_time = time.time()
top = Topology.from_openmm(pdbfile.topology, unique_molecules=[hexadecane])
print("--- Topology creation finished in %s seconds ---" % (time.time() - start_time))
ff = ForceField('test_forcefields/smirnoff99Frosst.offxml')

--- Topology creation finished in 1585.9495000839233 seconds ---

start_time = time.time()

system = ff.create_openmm_system(top)
print("--- Ssytem creation finished in %s seconds ---" % (time.time() - start_time))

--- Ssytem creation finished in 57.06095290184021 seconds ---

output.pdb.gz

Thoughts

I suspect that this slowdown occurs in the graph matching stage of topology creation. Hexadecane has a huge number of self-symmetries. While the function only takes the first molecular topology match that's found, it may be generating all of them unnecessarily.

Relevant code is here: https://github.com/openforcefield/openforcefield/blob/master/openforcefield/topology/topology.py#L1635-L1664

We should poke around NetworkX to see if there's a faster algorithm that this code could use.

[jaimergp]

I had to do something similar (identify copies of small molecules in a system) and found that "serializing" all connected molecules to canonical SMILES (as provided by rdkit) and comparing those was enough for identification purposes. Would that be enough in this case or am I missing something?

[jchodera]

It looks like you also do the expensive graph matching for all molecules, then try to correct for having different mappings afterwards. We really need to do some fast hashing / caching to avoid this expense.

What if we used an algorithm like this:

• For each Topology molecule, construct a hashable immutable object (like a tuple) that encodes the sequence of elements and bond indices within the molecule:

( (element_1, element_2, ..., element_N), ((bond1_atomindex1, bond1_atomindex2), ..., (bondN_atomindex1, bond1_atomindex2))

• Check a dict to see if we have already cached a matched isomorphism to a unique molecule. If not, perform the graph isomorphism match and cache it in the dict.
• Use the cached mapping to the unique molecule

If the Topology object consists of lots of repeated instances of the same molecule, this procedure would run the expensive graph isomorphism code exactly once, which should speed things up by ~1000-fold here.

[jchodera]

In the case where many candidate molecules are provided (which is not this case, but could be a future use case), we would also want to speed up the determination of which unique molecule the Topology molecule maps to. You should also be able to use a fast hashing there based on, for example, the Hill formula (e.g. C16H34), which should narrow down the list of possible graph isomorphisms to evaluate.

From there, imperfect matches can also be rapidly rejected by first calling could_be_isomorphic() before attempting to compute complete graph isomorphisms.

Later on, we'll have to think about optimizing this for the case where one of the molecules in the system is a biological macromolecule.

[SimonBoothroyd]

In the past when solving this problem I've always done very similar to what @jchodera suggests and find that in general it works well.

The only slight difference is that before even considering the unique_molecules, I would suggest going through all of the molecules in the system, building a hash for them, and then construct a dict of molecule # to a list of molecule indices which share the hash. The end result should in principle should be a dictionary with very few keys which should correspond to the number of components in the system.

You would then do the graph isomorphism matching on the first molecule in each of the lists in the dictionary map. In doing this you replace a lot of graph matching with simple hashing, and can safely only consider one molecule of each type when doing the matching (provided the hash is sensible).

The added benefit is that you can do some nice sanity checks and provide some feedback to the user. Namely, if you find that two of your hashes map to the same unique molecule something most likely is going wrong (e.g think of a pure system where for some reason one of the molecules has a slightly different atom ordering than the others in the system), and additionally if a unique molecule has not been matched to anything in the system. This could also be used to easily check / enforce that all molecules of the same type have the same atom ordering and are ordered contiguously which in turn could yield performance improvements and code simplification when parameterising the system.

This could also require (for openmm systems I think?) code to turn lists of atoms and bonds into individual molecules, but this isn't too bad to do and is pretty quick to run (you can use a flood-fill like algorithm).

The pseudocode would be something like:

all_molecules = molecules_from_openmm_topology(topology)


def hash_molecule(molecule):

    atom_tuples = (atom.element for atom in molecule.atoms)

    bond_list = [(min(bond.index1, bond.index2),
                  max(bond.index1, bond.index2), bond.order) for bond in molecule.bonds]

    bond_list.sort()

    bond_tuples = tuple(bond_list)

    molecule_tuple = (atom_tuples, bond_tuples)
    return hash(molecule_tuple)


distinguishable_molecules = {}

for index, molecule in enumerate(all_molecules):

    molecule_hash = hash_molecule(molecule)

    if molecule_hash not in distinguishable_molecules:
        distinguishable_molecules[molecule_hash] = []

    distinguishable_molecules[molecule_hash].append(index)

hash_to_unique_molecule = {}
unique_molecule_to_hash = {}

for molecule_hash in distinguishable_molecules:

    molecule = distinguishable_molecules[molecule_hash][0]

    for index, unique_molecule in enumerate(unique_molecules):

        is_match = check_isomorphism(molecule, unique_molecule)
        
        if not is_match:
            continue

        if molecule_hash in hash_to_unique_molecule:
            raise Exception('Molecule in system matches multiple unique mols.')

        hash_to_unique_molecule[molecule_hash] = unique_molecule
        
        if index in unique_molecule_to_hash:
            
            logging.warning('Multiple molecules of the same type have a different '
                            'atom / bond ordering - this is most likely unexpected '
                            'behaviour...')
            
            continue
            
        unique_molecule_to_hash[index] = molecule_hash

[j-wags]

I had to do something similar (identify copies of small molecules in a system) and found that "serializing" all connected molecules to canonical SMILES (as provided by rdkit) and comparing those was enough for identification purposes. Would that be enough in this case or am I missing something?

@jaimergp This would be sufficient for identification, but unfortunately we need to get an atom index mapping from this process.

[j-wags]

@jaimergp @jchodera @SimonBoothroyd Thanks for the really thoughtful feedback here. I agree that these ideas should net us a close to 1000x speedup in this case. This is super helpful :-)

[j-wags]

@SimonBoothroyd

This could also require (for openmm systems I think?) code to turn lists of atoms and bonds into individual molecules, but this isn't too bad to do and is pretty quick to run (you can use a flood-fill like algorithm).

Alas! A critical facet of our implementation is that we don't require detailed molecule info (bond orders and stereochemistry) from OpenMM. If we had a reliable source of these, we wouldn't need unique_molecules at all. The pseudocode is still really helpful though.

[SimonBoothroyd]

So just to clarify - when I say into individual molecules, I mean each atom needs to be assigned a parent molecule index (i.e atoms [0-19] belong are assigned a molecule index of 0, atoms [20-39] are assigned a molecule index of 1) as opposed to we need to create actual Molecule objects. i.e I really mean we just need to know which groups of atoms are part of the same molecule for the hashing approach to work. As a very rough sketch, something more like:

namedtuple('OpenMMMolecule', ['atom_indices', 'bond_indices'])

If the bond order is there, awesome, if not, a bond order of None still works in the hash!

[j-wags]

@jthorton and I have been discussing consolidating/deduplicating molecule identity operations. I think that making two highly-performant methods can cover all of our needs:

• staticmethod FrozenMolecule.is_isomorphic, which returns bool
• staticmethod FrozenMolecule.get_atom_map, which returns dict{int: int}

The arguments for both methods can be (mol1, mol2, atomic_number=True, atom_index=True, atom_is_aromatic=True, atom_stereo=True, bond_order=True, bond_is_aromatic=True, bond_stereo=True)

mol1 and mol2 can be FrozenMolecule, TopologyMolecule, or nx.Graph (any others?)

If the user specifies bond_order=True, but the nx.Graph that they provide doesn't have bond orders, we should raise a ValueError. We might consider having several subclasses of ValueError for the different deficiencies we could find.

I'm in favor of keeping these two functions separate, because I don't like when one function can return different types (bool vs. dict). This tends to confuse static code analysis and users.

The question remains of "can we do this in a way to efficiently handle batch operations?". Caching data on the Molecule objects might speed this up a bit. I'll need to think more about whether these functions will be compatible with implementing something like @SimonBoothroyd's algorithm above.

[SimonBoothroyd]

I'm in favor of keeping these two functions separate, because I don't like when one function can return different types (bool vs. dict). This tends to confuse static code analysis and users.

Or you could just have a staticmethod FrozenMolecule.are_isomorphic with a signature like:

are_isomorphic, atom_map = FrozenMolecule.are_isomorphic(mol1, mol2, ...)

where are_isomorphic: bool, atom_map: Optional[Dict[int, int]], which would probably make sense given that get_atom_map would likely need to call is_isomorphic anyway, and both methods would likely have a large code overlap so the additional cost to are_isomorphic would be minimal.

I'll need to think more about whether these functions will be compatible with implementing something like @SimonBoothroyd's algorithm above.

They would - they would just replace the check_isomorphism call.

Caching data on the Molecule objects might speed this up a bit

I would advise against caching unless essential - the required bookkeeping associated with caching is never fun, and avoiding caching would ensure that the methods to work well with Molecule without modification.

[j-wags]

Ahh, yes. I agree -- We should make a single function that returns are_isomorphic: bool and atom_map: dict{int:int} or None. The kwarg return_atom_map would determine whether atom_map is returned as a dict, otherwise it'll be None. So the signature would look like:

are_isomorphic, atom_map = FrozenMolecule.are_isomorphic(mol1, mol2, return_atom_map=True, atomic_number=True, atom_index=True, atom_is_aromatic=True, atom_stereo=True, bond_order=True, bond_is_aromatic=True, bond_stereo=True)

(Reading through our DM conversation, I now see this may be what @jthorton was suggesting all along)

I agree that caching can be a pain to manage. On the other hand, the most efficient isomorphism checks in some cases overlap with other OFFTK functionality -- eg. SMILES are effectively molecule hashes for certain comparisons. Specifically, are_isomorphic(mol1, mol2, atom_map=False, atom_index=False) is equivalent to mol1.to_smiles()==mol2.to_smiles(), and we already cache SMILES. Further, the same query with atom_index=True could be performed by comparing the Canonical, isomeric, explicit hydrogen, mapped CMILES, and as we more tightly integrate the OFF ecosystem, users may find it helpful to have CMILES cached on OFFMols. That said, not all hashes will be useful, and maybe it's not useful to store anything past those two.

One of the other concerns is that the proposed function will only do 1 vs. 1 comparisons, meaning that matching a dataset of N molecules to M unique mols is an O(N*M) operation. However, if we could extract the hashes, then we could make a hash table for either M or N, and complete the matching in O(M+N) time using John/Simon's algorithm. Since we already have the algorithm in mind, and it is a significant "big-O-notation" speed-up for dataset matching, maybe we should have FrozenMolecule.are_isomorphic be able to accept either mol1 or mol2 as iterables (eg. lists, tuples) of molecules, and return the M x N array of bool matches and another MxN array of dicts/Nones