Ensure downloader selects the correct level of theory (#47)

openmm · Sep 29, 2022 · 0659915 · 0659915
1 parent 927fbae
commit 0659915
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/downloader/downloader.py b/downloader/downloader.py
@@ -77,8 +77,11 @@ def compute_reference_energy(smiles):
     print('Processing', subset)
     ds = client.get_collection('Dataset', subset)
     all_molecules = ds.get_molecules()
-    spec = ds.list_records().iloc[0].to_dict()
-    recs = ds.get_records(method=spec['method'], basis=spec['basis'], program=spec['program'], keywords=spec['keywords'])
+    for row in ds.list_records().iloc:
+        spec = row.to_dict()
+        if spec['method'] == 'wb97m-d3bj':
+            recs = ds.get_records(method=spec['method'], basis=spec['basis'], program=spec['program'], keywords=spec['keywords'])
+            break
     recs_by_name = defaultdict(list)
     mols_by_name = defaultdict(list)
     for i in range(len(recs)):