Is your feature request related to a problem? Please describe.
Importing just about anything from the toolkit results in heavy modules such as pytorch_lightning also being imported. Here's a visualization from an one-liner (python -X importtime -c "from openff.toolkit import ForceField" 2> import.log && tuna import.log) but the results are similar if you pull in other classes:
Most use cases with force fields will indeed require all of this to be imported, but it would be convenient to have this happen only when it's really needed (force field parametrization time) not beforehand (just importing core classes).
My reading of the code tells me the chain of events is
- Does anything pull in something in
openff.toolkit.topology?
- If so, we're going to be importing (all of)
openff.toolkit.utils
- Included is all of the toolkit wrappers and and instantiation of
GLOBAL_TOOLKIT_REGISTRY
- Providing
GLOBAL_TOOLKIT_REGISTRY requires instantiating a ToolkitRegistry and all of its constituent ToolkitWrappers, which includes NAGLToolkitWrapper
- Naturally, making a
NAGLToolkitWrapper requires importing from openff.nagl one way or another, and this will hit PyTorch-based code paths
Describe the solution you'd like
It would be great to defer the PyTorch imports until GNN-based inference actually happens.
Describe alternatives you've considered
Additional context
This is at bottom the same as openforcefield/openff-nagl#293 but the user-facing expectations here and there are, I believe, a little different
Is your feature request related to a problem? Please describe.
Importing just about anything from the toolkit results in heavy modules such as
pytorch_lightningalso being imported. Here's a visualization from an one-liner (python -X importtime -c "from openff.toolkit import ForceField" 2> import.log && tuna import.log) but the results are similar if you pull in other classes:Most use cases with force fields will indeed require all of this to be imported, but it would be convenient to have this happen only when it's really needed (force field parametrization time) not beforehand (just importing core classes).
My reading of the code tells me the chain of events is
openff.toolkit.topology?openff.toolkit.utilsGLOBAL_TOOLKIT_REGISTRYGLOBAL_TOOLKIT_REGISTRYrequires instantiating aToolkitRegistryand all of its constituentToolkitWrappers, which includesNAGLToolkitWrapperNAGLToolkitWrapperrequires importing fromopenff.naglone way or another, and this will hit PyTorch-based code pathsDescribe the solution you'd like
It would be great to defer the PyTorch imports until GNN-based inference actually happens.
Describe alternatives you've considered
Additional context
This is at bottom the same as openforcefield/openff-nagl#293 but the user-facing expectations here and there are, I believe, a little different