Improved compute

molgenevalmetric.py

@@ -223,7 +223,7 @@ def average_sascore(gen, n_jobs=1):
 
 def average_agg_tanimoto(stock_vecs, gen_vecs,
                          batch_size=5000, agg='max',
-                         device='cpu', p=1):
+                         device=None, p=1):
     """
     Calculates the average aggregate Tanimoto similarity between two sets of molecule fingerprints.
 
@@ -232,7 +232,7 @@ def average_agg_tanimoto(stock_vecs, gen_vecs,
     - gen_vecs (numpy array): Fingerprint vectors for the generated molecule set.
     - batch_size (int): The size of batches to process similarities (reduces memory usage).
     - agg (str): Aggregation method, either 'max' or 'mean'.
-    - device (str): The computation device ('cpu' or 'cuda:0', etc.).
+    - device (str or None): The computation device ('cpu' or 'cuda:0', etc.). If None, automatically detect.
     - p (float): The power for averaging, used in generalized mean calculation.
 
     Returns:
@@ -240,6 +240,11 @@ def average_agg_tanimoto(stock_vecs, gen_vecs,
     """
 
     assert agg in ['max', 'mean'], "Can aggregate only max or mean"
+    
+    # Automatically detect device if not provided
+    if device is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
     agg_tanimoto = np.zeros(len(gen_vecs))
     total = np.zeros(len(gen_vecs))
     for j in range(0, stock_vecs.shape[0], batch_size):
@@ -553,24 +558,34 @@ class molgenevalmetric(evaluate.Metric):
         )
 
     def _compute(self, gensmi, trainsmi):
-
         metrics = {}
-        metrics['Novelty'] = novelty(gen = gensmi, train = trainsmi)
-        metrics['Valid'] = fraction_valid(gen=gensmi)
-        metrics['Unique'] = fraction_unique(gen=gensmi)
-        metrics['IntDiv'] = internal_diversity(gen=gensmi)
-        metrics['FCD'] = fcd_metric(gen = gensmi, train = trainsmi)
-        metrics['QED'] = qed_metric(gen=gensmi)
-        metrics['LogP'] = logP_metric(gen=gensmi)
-        metrics['Penalized LogP'] = penalized_logp(gen=gensmi)
-        metrics['SA'] = average_sascore(gen=gensmi)
-        metrics['SCScore'] = synthetic_complexity_score(gen=gensmi)
-        metrics['SYBA'] = SYBAscore(gen=gensmi)
-        # metrics['Oracles'] = oracles(gen = gensmi, train = trainsmi)
+        metric_functions = {
+            'Novelty': lambda: novelty(gen=gensmi, train=trainsmi),
+            'Valid': lambda: fraction_valid(gen=gensmi),
+            'Unique': lambda: fraction_unique(gen=gensmi),
+            'IntDiv': lambda: internal_diversity(gen=gensmi),
+            'FCD': lambda: fcd_metric(gen=gensmi, train=trainsmi),
+            'QED': lambda: qed_metric(gen=gensmi),
+            'LogP': lambda: logP_metric(gen=gensmi),
+            'Penalized LogP': lambda: penalized_logp(gen=gensmi),
+            'SA': lambda: average_sascore(gen=gensmi),
+            'SCScore': lambda: synthetic_complexity_score(gen=gensmi),
+            'SYBA': lambda: SYBAscore(gen=gensmi),
+            # 'Oracles': lambda: oracles(gen=gensmi, train=trainsmi)
+        }
+
+        for metric_name, compute_func in metric_functions.items():
+            print(f"Computing {metric_name}...")
+            try:
+                metrics[metric_name] = compute_func()
+            except Exception as e:
+                print(f"Error computing {metric_name}: {e}")
+                metrics[metric_name] = 0
 
         return metrics
 
 
+
 # def get_n_rings(mol):
 #     """
 #     Computes the number of rings in a molecule