Speed up TorsionDriveResultCollection.to_records by batching requests (…

…#270) * 4x speedup in TorsionDriveResultCollection by batching requests * include final_molecule for optimization too * update release notes --------- Co-authored-by: Jeffrey Wagner <[email protected]>
openforcefield · Mar 22, 2024 · 223721d · 223721d
1 parent 39c9664
commit 223721d
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 7 deletions.
diff --git a/docs/releasehistory.md b/docs/releasehistory.md
@@ -17,6 +17,10 @@ Releases are given with dates in DD-MM-YYYY format.
 * [PR #260:] Fixes a bug where adding different tautomers of the same molecule to a ComponentResult would raise an error. Also fixes a case where ComponentResult.add_molecule would fail to return a bool  (#255) [@mattwthompson]
 * [PR #268:] Fix broken star imports (#268) [@mattwthompson]
 
+### Performance Improvements
+
+* [PR #270:] Speed up `TorsionDriveResultCollection.to_records` by batching requests [@ntBre]
+
 <!--## Version / Date DD-MM-YYYY -->
 ## 0.50.2 / 24-01-2024
 
@@ -106,6 +110,7 @@ For more information on this release, see https://github.com/openforcefield/open
 [PR #257:]: https://github.com/openforcefield/openff-qcsubmit/pull/257
 [PR #260:]: https://github.com/openforcefield/openff-qcsubmit/pull/260
 [PR #268:]: https://github.com/openforcefield/openff-qcsubmit/pull/268
+[PR #270:]: https://github.com/openforcefield/openff-qcsubmit/pull/270
 
 [@jthorton]: https://github.com/jthorton
 [@dotsdl]: https://github.com/dotsdl
@@ -114,3 +119,4 @@ For more information on this release, see https://github.com/openforcefield/open
 [@chapincavender]: https://github.com/chapincavender
 [@j-wags]: https://github.com/j-wags
 [@pavankum]: https://github.com/pavankum
+[@ntBre]: https://github.com/ntBre
diff --git a/openff/qcsubmit/results/results.py b/openff/qcsubmit/results/results.py
@@ -538,7 +538,7 @@ def to_records(self) -> List[Tuple[OptimizationRecord, Molecule]]:
             rec_ids = [result.record_id for result in results]
             # Do one big request to save time
             opt_records = client.get_optimizations(
-                rec_ids, include=["initial_molecule"]
+                rec_ids, include=["initial_molecule", "final_molecule"]
             )
             # Sort out which records from the request line up with which results
             opt_rec_id_to_result = dict()
@@ -804,9 +804,13 @@ def to_records(self) -> List[Tuple[TorsiondriveRecord, Molecule]]:
         for client_address, records in self.entries.items():
             client = PortalClient(client_address)
 
-            for record in records:
-                rec = client.get_torsiondrives(record.record_id)
-
+            # retrieve all torsiondrives at once, including their
+            # minimum_optimizations
+            record_ids = [r.record_id for r in records]
+            torsion_drive_records = client.get_torsiondrives(
+                record_ids, include=["minimum_optimizations"]
+            )
+            for record, rec in zip(records, torsion_drive_records):
                 # OpenFF molecule
                 try:
                     molecule: Molecule = Molecule.from_mapped_smiles(
@@ -819,10 +823,16 @@ def to_records(self) -> List[Tuple[TorsiondriveRecord, Molecule]]:
                     )
                     continue
 
+                # retrieve all minimum_optimizations at once
+                opt_ids = [v.id for v in rec.minimum_optimizations.values()]
+                opt_records = client.get_optimizations(opt_ids)
+                # retrieve all final_molecules at once
+                opt_final_molecule_ids = [r.final_molecule_id for r in opt_records]
+                opt_final_molecules = client.get_molecules(opt_final_molecule_ids)
                 # Map of torsion drive keys to minimum optimization
-                qc_grid_molecules = [
-                    (k, v.final_molecule) for k, v in rec.minimum_optimizations.items()
-                ]
+                qc_grid_molecules = list(
+                    zip(rec.minimum_optimizations.keys(), opt_final_molecules)
+                )
 
                 # order the ids so the conformers follow the torsiondrive scan range
                 # x[0] is the torsiondrive key, ie Tuple[float]