fix drop duplicates

NREL · Sep 22, 2021 · 80792b8 · 80792b8
1 parent c466a1b
commit 80792b8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/alfabet/prediction.py b/alfabet/prediction.py
@@ -64,8 +64,7 @@ def predict_bdes(smiles, draw=False):
                             on=['molecule', 'bond_index'], how='left')
 
     # Drop duplicate entries and sort from weakest to strongest
-    frag_df = frag_df.sort_values('bde_pred').drop_duplicates(
-        ['fragment1', 'fragment2']).reset_index(drop=True)
+    frag_df = frag_df.sort_values('bde_pred').reset_index(drop=True)
 
     # Draw SVGs
     if draw:

diff --git a/tests/test_model.py b/tests/test_model.py
@@ -1,6 +1,8 @@
 import numpy as np
+
 from alfabet import model
 
+
 def test_predict():
     results = model.predict(['CC', 'NCCO', 'CF', 'B'], verbose=False)
 
@@ -16,4 +18,12 @@ def test_predict():
 
     np.testing.assert_allclose(
         results[results.molecule == 'NCCO'].bde_pred,
-        [90.0, 82.1, 98.2, 99.3, 92.1, 92.5, 105.2], atol=1., rtol=.05)
+        [90.0, 82.1, 98.2, 99.3, 92.1, 92.5, 105.2], atol=1., rtol=.05)
+
+
+def test_duplicates():
+    results = model.predict(['c1ccccc1'], verbose=False, drop_duplicates=True)
+    assert len(results) == 1
+
+    results = model.predict(['c1ccccc1'], verbose=False, drop_duplicates=False)
+    assert len(results) == 6