Improvements to record downloading and docstrings (#352)

* Improvements to record downloading and docstrings

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Apply corrections from code review.

Co-authored-by: Lily Wang <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused import

* restore support for single drivers outside iterable

* update releasehistory

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lily Wang <[email protected]>
Co-authored-by: Jeff Wagner <[email protected]>

Author: 4 people

docs/releasehistory.md

@@ -11,6 +11,16 @@ Releases are given with dates in DD-MM-YYYY format.
 
 <!--## Version / Date DD-MM-YYYY -->
 
+## 0.57.0 / 10-02-2025
+
+### API extensions
+* [PR #352:] Add `include` argument to `to_records()` to enable users to select which fields to download (and hint to user that not all fields are downloaded by default!)
+
+### Documentation updates
+* [PR #352:] Revise docstrings for to_records() to better reflect their purpose
+* [PR #352:] Add a note to OptimizationResultCollection.to_records() docstring to explain that the user may want to call .to_basic_result_collection() first
+* [PR #352:] Revise docstring of to_basic_result_collection() with more details
+
 ## 0.56.0 / 05-19-2025
 
 ### Bugfixes

openff/qcsubmit/results/results.py

@@ -192,9 +192,17 @@ def _validate_record_types(
             )
 
     @abc.abstractmethod
-    def to_records(self) -> List[Tuple[BaseRecord, Molecule]]:
-        """Returns the native QCPortal record objects for each of the records referenced
-        in this collection along with a corresponding OpenFF molecule object.
+    def to_records(self, include: Iterable[str]) -> List[Tuple[BaseRecord, Molecule]]:
+        """Download all records referenced in this collection from QCFractal.
+
+        Returns the native QCPortal record objects for each of the records
+        referenced in this collection along with a corresponding OpenFF Molecule
+        object.
+
+        Parameters
+        ==========
+        include
+            The fields to download when the record is collected.
         """
         raise NotImplementedError()
 
@@ -388,22 +396,36 @@ def from_server(
             spec_name,
         )
 
-    def to_records(self) -> List[Tuple[SinglepointRecord, Molecule]]:
-        """Returns the native QCPortal record objects for each of the records referenced
-        in this collection along with a corresponding OpenFF molecule object.
+    def to_records(
+        self, include: Iterable[str] = ("molecule",)
+    ) -> List[Tuple[SinglepointRecord, Molecule]]:
+        """Download all records referenced in this collection from QCFractal.
+
+        Returns the native QCPortal record objects for each of the records
+        referenced in this collection along with a corresponding OpenFF Molecule
+        object.
 
         Each molecule will contain the conformer referenced by the record.
+
+        Parameters
+        ==========
+        include
+            The fields to download when the record is collected. The
+            ``"molecule"`` field is always downloaded.
         """
         from openff.qcsubmit.utils.utils import _default_portal_client
 
         records_and_molecules = []
 
+        if "molecule" not in include:
+            include = ("molecule", *include)
+
         for client_address, records in self.entries.items():
             client = _default_portal_client(client_address)
 
             # TODO - batching/chunking (maybe in portal?)
             for record in records:
-                rec = client.get_singlepoints(record.record_id, include=["molecule"])
+                rec = client.get_singlepoints(record.record_id, include=include)
 
                 # OpenFF molecule
                 try:
@@ -540,25 +562,45 @@ def from_server(
             spec_name,
         )
 
-    def to_records(self) -> List[Tuple[OptimizationRecord, Molecule]]:
-        """Returns the native QCPortal record objects for each of the records referenced
-        in this collection along with a corresponding OpenFF molecule object.
+    def to_records(
+        self,
+        include: Iterable[str] = ("initial_molecule", "final_molecule"),
+    ) -> List[Tuple[OptimizationRecord, Molecule]]:
+        """Download all records referenced in this collection from QCFractal.
+
+        Returns the native QCPortal record objects for each of the records
+        referenced in this collection along with a corresponding OpenFF Molecule
+        object.
 
         Each molecule will contain the minimum energy conformer referenced by the
         record.
+
+        Parameters
+        ==========
+        include
+            The fields to download when the record is collected. The
+            ``"final_molecule"`` field is always downloaded.
+
+        Notes
+        =====
+        By default, this function does not download ``"trajectory"`` field,
+        and so the resulting records do not include any gradients or forces. To
+        download a complete record of the minimum energy conformer's calculation,
+        including forces, see :py:meth:`to_basic_result_collection`.
         """
         from openff.qcsubmit.utils.utils import _default_portal_client
 
         records_and_molecules = []
 
+        if "final_molecule" not in include:
+            include = ("final_molecule", *include)
+
         for client_address, results in self.entries.items():
             client = _default_portal_client(client_address)
 
             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", "final_molecule"]
-            )
+            opt_records = client.get_optimizations(rec_ids, include=include)
             # Sort out which records from the request line up with which results
             opt_rec_id_to_result = dict()
             for result in results:
@@ -593,18 +635,34 @@ def to_records(self) -> List[Tuple[OptimizationRecord, Molecule]]:
 
         return records_and_molecules
 
-    def to_basic_result_collection(self, driver) -> BasicResultCollection:
-        """Returns a basic results collection which references results records which
+    def to_basic_result_collection(
+        self,
+        driver: SinglepointDriver | Iterable[SinglepointDriver] | None = None,
+    ) -> BasicResultCollection:
+        """
+        Get a collection of the single point results from the end of each optimization.
+
+        Returns a basic results collection which references results records which
         were created from the *final* structure of one of the optimizations in this
         collection, and used the same program, method, and basis as the parent
         optimization record.
 
-        Returns:
-            The results collection referencing records created from the final optimized
-            structures referenced by this collection.
+        Parameters
+        ==========
+        driver
+            Return only those records whose driver is in this list. If omitted
+            or ``None``, include all drivers.
+
+        Returns
+        =======
+        The results collection referencing records created from the final
+        optimized structures referenced by this collection.
         """
         from openff.qcsubmit.utils.utils import _default_portal_client
 
+        # If driver is None, set it to all drivers
+        driver = tuple(SinglepointDriver) if driver is None else driver
+
         records_and_molecules = self.to_records()
 
         result_records = defaultdict(list)
@@ -812,25 +870,40 @@ def from_server(
             spec_name,
         )
 
-    def to_records(self) -> List[Tuple[TorsiondriveRecord, Molecule]]:
-        """Returns the native QCPortal record objects for each of the records referenced
-        in this collection along with a corresponding OpenFF molecule object.
+    def to_records(
+        self,
+        include: Iterable[str] = ("minimum_optimizations",),
+    ) -> List[Tuple[TorsiondriveRecord, Molecule]]:
+        """Download all records referenced in this collection from QCFractal.
+
+        Returns the native QCPortal record objects for each of the records
+        referenced in this collection along with a corresponding OpenFF Molecule
+        object.
 
         Each molecule will contain the minimum energy conformer referenced by the
         record.
+
+        Parameters
+        ==========
+        include
+            The fields to download when the record is collected. The
+            ``"minimum_optimizations"`` field is always downloaded.
         """
         from openff.qcsubmit.utils.utils import _default_portal_client
 
         records_and_molecules = []
 
+        if "minimum_optimizations" not in include:
+            include = ("minimum_optimizations", *include)
+
         for client_address, records in self.entries.items():
             client = _default_portal_client(client_address)
 
             # 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"]
+                record_ids, include=include
             )
             for record, rec in zip(records, torsion_drive_records):
                 # OpenFF molecule