Merge pull request #11 from mattwthompson/sort-qcarchive-id

Sort getters by QCArchive ID
openforcefield · Dec 5, 2023 · 1b661cf · 1b661cf
2 parents 7dd24b6 + 3b93dcb
commit 1b661cf
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git a/ibstore/_store.py b/ibstore/_store.py
@@ -224,9 +224,11 @@ def get_qcarchive_ids_by_molecule_id(self, id: int) -> list[str]:
                 qcarchive_id
                 for (qcarchive_id,) in db.db.query(DBQMConformerRecord.qcarchive_id)
                 .filter_by(parent_id=id)
+                .order_by(DBQMConformerRecord.qcarchive_id)
                 .all()
             ]
 
+    # TODO: if this can take a list of ids, should it sort by QCArchive ID
     def get_molecule_id_by_qcarchive_id(self, id: str) -> int:
         with self._get_session() as db:
             return [
@@ -287,6 +289,7 @@ def get_qm_energies_by_molecule_id(self, id: int) -> list[float]:
                 energy
                 for (energy,) in db.db.query(DBQMConformerRecord.energy)
                 .filter_by(parent_id=id)
+                .order_by(DBQMConformerRecord.qcarchive_id)
                 .all()
             ]
 
@@ -302,6 +305,7 @@ def get_mm_energies_by_molecule_id(
                 for (energy,) in db.db.query(DBMMConformerRecord.energy)
                 .filter_by(parent_id=id)
                 .filter_by(force_field=force_field)
+                .order_by(DBQMConformerRecord.qcarchive_id)
                 .all()
             ]
 

diff --git a/ibstore/_tests/unit_tests/test_store.py b/ibstore/_tests/unit_tests/test_store.py
@@ -1,7 +1,10 @@
+import json
+import random
 import tempfile
 
 import numpy
 import pytest
+from openff.qcsubmit.results import OptimizationResultCollection
 from openff.utilities import get_data_file_path, temporary_cd
 
 from ibstore._store import MoleculeStore
@@ -57,6 +60,28 @@ def test_get_molecule_id_by_qcarchive_id(small_store):
     assert small_store.get_molecule_id_by_qcarchive_id(qcarchive_id) == molecule_id
 
 
+def test_molecules_sorted_by_qcarchive_id():
+    raw_ch = json.load(
+        open(get_data_file_path("_tests/data/01-processed-qm-ch.json", "ibstore")),
+    )
+
+    random.shuffle(raw_ch["entries"]["https://api.qcarchive.molssi.org:443/"])
+
+    with tempfile.NamedTemporaryFile(mode="w+") as file:
+        json.dump(raw_ch, file)
+        file.flush()
+
+        store = MoleculeStore.from_qcsubmit_collection(
+            OptimizationResultCollection.parse_file(file.name),
+            database_name=tempfile.NamedTemporaryFile(suffix=".sqlite").name,
+        )
+
+        qcarchive_ids = store.get_qcarchive_ids_by_molecule_id(40)
+
+    for index, id in enumerate(qcarchive_ids[:-1]):
+        assert id < qcarchive_ids[index + 1]
+
+
 def test_get_conformers(small_store):
     force_field = "openff-2.0.0"
     molecule_id = 40