Cache some variables during MSA featurization #288
Conversation
amorehead
commented
Sep 23, 2024
amorehead
commented
- Caches some variables during MSA featurization
| for idx, (chemtype, residue_index) in enumerate( | ||
| zip(chain_chemtype, chain_residue_index) | ||
| ): | ||
| is_polymer = chemtype < ligand_chemtype_index | ||
| is_ligand = not is_polymer | ||
|
|
||
| chem_residue_constants = get_residue_constants(res_chem_index=chemtype) | ||
| if chemtype not in chemtype_constants_cache: | ||
| chemtype_constants_cache[chemtype] = get_residue_constants(res_chem_index=chemtype) |
There was a problem hiding this comment.
is get_residue_constants an expensive function?
There was a problem hiding this comment.
I need to run some more experiments, but caching its execution for each token in a sequence seems to have dropped the runtime of pdb_input_to_atom_input from minutes to seconds. Fingers crossed this pattern holds for other structures
There was a problem hiding this comment.
In principle, it should be a O(1) complexity function, since it does a simple lookup based on integers or strings and returns a static Python Module class instance. However, the make_msa_features function is in total a O(C * M * S) complexity function, where C is the number of chains in a structure (e.g., 3), M is the number of sequences in an MSA (e.g., 16k), and S is the number of tokens in each MSA sequence (e.g., 200).
