Merge pull request #20 from aakankschit/feature/custom-evaluator

Add CustomEvaluator implementation

Author: aakankschit

src/TACTICS/thompson_sampling/config.py

@@ -30,6 +30,7 @@
     ROCSEvaluatorConfig,
     FredEvaluatorConfig,
     MLClassifierEvaluatorConfig,
+    CustomEvaluatorConfig,
 )
 
 if TYPE_CHECKING:
@@ -59,6 +60,7 @@
     ROCSEvaluatorConfig,
     FredEvaluatorConfig,
     MLClassifierEvaluatorConfig,
+    CustomEvaluatorConfig,
 ]
 
 

src/TACTICS/thompson_sampling/core/__init__.py

@@ -13,7 +13,8 @@
     FredEvaluator,
     FPEvaluator,
     MWEvaluator,
-    MLClassifierEvaluator
+    MLClassifierEvaluator,
+    CustomEvaluator,
 )
 
 __all__ = [
@@ -26,4 +27,5 @@
     'FPEvaluator',
     'MWEvaluator',
     'MLClassifierEvaluator',
+    'CustomEvaluator',
 ] 

src/TACTICS/thompson_sampling/core/evaluator_config.py

@@ -1,7 +1,7 @@
 """Pydantic configuration models for evaluators."""
 
 from pydantic import BaseModel, Field
-from typing import Literal, Optional
+from typing import Literal, Optional, Callable, Any
 
 
 class LookupEvaluatorConfig(BaseModel):
@@ -88,3 +88,17 @@ class MLClassifierEvaluatorConfig(BaseModel):
 
     evaluator_type: Literal["ml_classifier"] = "ml_classifier"
     model_filename: str = Field(..., description="Path to trained model file (joblib pickle)")
+
+class CustomEvaluatorConfig(BaseModel):
+    """
+    Configuration for CustomEvaluator.
+
+    Uses a user-provided callable scoring function that accepts an RDKit Mol
+    and returns a float.
+    """
+
+    evaluator_type: Literal["custom"] = "custom"
+    scoring_function: Callable[[Any], float] = Field(
+        ...,
+        description="Callable that accepts an RDKit Mol and returns a float score"
+    )

src/TACTICS/thompson_sampling/core/evaluators.py

@@ -234,6 +234,42 @@ def evaluate(self, mol):
             score = float(oechem.OEGetSDData(docked_mol, sd_tag))
         return score
 
+class CustomEvaluator(Evaluator):
+    """An evaluator class that uses a user-provided custom scoring function
+    """
+
+    def __init__(self, scoring_function):
+        """
+        :param scoring_function: callable that accepts an RDKit Mol and returns a float
+        """
+        self.scoring_function = scoring_function
+        self.num_evaluations = 0
+        self.score_cache = {}
+
+    @property
+    def counter(self):
+        return self.num_evaluations
+
+    def evaluate(self, mol):
+        """Evaluate a molecule using the user-defined scoring function
+        :param mol: Input RDKit molecule
+        :return: Float score from the custom function, np.nan on failure
+        """      
+        self.num_evaluations += 1      
+        
+        # Look up to see if we already processed this molecule
+        smi = Chem.MolToSmiles(mol)
+        cached_score = self.score_cache.get(smi)
+        if cached_score is not None:
+            return cached_score
+
+        try:
+            score = float(self.scoring_function(mol))
+        except Exception:
+            score = np.nan
+
+        self.score_cache[smi] = score
+        return score
 
 def generate_confs(mol, max_confs):
     """Generate conformers with Omega

src/TACTICS/thompson_sampling/factories.py

@@ -41,7 +41,8 @@
     MWEvaluatorConfig,
     ROCSEvaluatorConfig,
     FredEvaluatorConfig,
-    MLClassifierEvaluatorConfig
+    MLClassifierEvaluatorConfig,
+    CustomEvaluatorConfig
 )
 from .core.evaluators import (
     LookupEvaluator,
@@ -51,7 +52,8 @@
     ROCSEvaluator,
     FredEvaluator,
     MLClassifierEvaluator,
-    Evaluator
+    Evaluator,
+    CustomEvaluator
 )
 
 
@@ -79,7 +81,8 @@
     MWEvaluatorConfig,
     ROCSEvaluatorConfig,
     FredEvaluatorConfig,
-    MLClassifierEvaluatorConfig
+    MLClassifierEvaluatorConfig,
+    CustomEvaluatorConfig
 ]
 
 
@@ -265,5 +268,9 @@ def create_evaluator(config: EvaluatorConfig) -> Evaluator:
         input_dict = {"model_filename": config.model_filename}
         return MLClassifierEvaluator(input_dict)
 
+    elif isinstance(config, CustomEvaluatorConfig):
+        scoring_function = config.scoring_function
+        return CustomEvaluator(scoring_function)
+
     else:
         raise ValueError(f"Unknown evaluator config type: {type(config)}")