Merge pull request #18 from coleygroup/pharm-prioritization

Conditional evaluation pipeline to support pharmacophore prioritization scoring

Author: kentoabeywardane

shepherd_score/conformer_generation.py

@@ -13,7 +13,7 @@
 from pathlib import Path
 from tqdm import tqdm
 import uuid
-from typing import Optional, List
+from typing import Optional, List, Union
 import contextlib
 import multiprocessing
 
@@ -308,7 +308,7 @@ def optimize_conformer_with_xtb(conformer: Chem.Mol,
                                 solvent: Optional[str] = None,
                                 num_cores: int = 1,
                                 charge: int = 0,
-                                temp_dir: str | Path = TMPDIR):
+                                temp_dir: Union[str, Path] = TMPDIR):
     """
     Use external calls to GFN2-XTB (command line) to optimize a conformer geometry.
 
@@ -395,7 +395,7 @@ def optimize_conformer_with_xtb_from_xyz_block(xyz_block: str,
                                                solvent: Optional[str] = None,
                                                num_cores: int = 1,
                                                charge: int = 0,
-                                               temp_dir: str | Path = TMPDIR):
+                                               temp_dir: Union[str, Path] = TMPDIR):
     """
     Use external calls to GFN2-XTB (command line) to optimize coordinates from an xyz block.
 
@@ -483,7 +483,7 @@ def charges_from_single_point_conformer_with_xtb(conformer: Chem.Mol,
                                                  solvent: Optional[str] = None,
                                                  num_cores: int = 1,
                                                  charge: int = 0,
-                                                 temp_dir: str | Path = TMPDIR
+                                                 temp_dir: Union[str, Path] = TMPDIR
                                                  ):
     """
     Compute atomic partial charges from a single point xTB calculation of a provided conformer.
@@ -558,7 +558,7 @@ def single_point_xtb_from_xyz(xyz_block: str,
                               solvent: Optional[str] = None,
                               num_cores: int = 1,
                               charge: int = 0,
-                              temp_dir: str | Path = TMPDIR):
+                              temp_dir: Union[str, Path] = TMPDIR):
     """
     Compute energy and atomic partial charges from a single point xTB calculation.
 
@@ -648,7 +648,7 @@ def optimize_conformer_ensemble_with_xtb(conformers: List[Chem.Mol],
                                        num_processes: int = 1,
                                        num_workers: int = 1,
                                        charge: int = 0,
-                                       temp_dir: str | Path = TMPDIR,
+                                       temp_dir: Union[str, Path] = TMPDIR,
                                        verbose: bool = False):
     """
     GFN2-XTB geometry optimization for a list of conformers.
@@ -735,7 +735,7 @@ def generate_opt_conformers_xtb(smiles: str,
                                 MMFF_optimize: bool = True,
                                 num_processes: int = 1,
                                 num_workers: int = 1,
-                                temp_dir: str | Path = TMPDIR,
+                                temp_dir: Union[str, Path] = TMPDIR,
                                 verbose: bool = False,
                                 num_confs: int = 1000):
     """

shepherd_score/evaluations/evaluate/_pipeline_eval_single.py

@@ -53,10 +53,10 @@ def _create_failed_result(i: int, error_msg: str) -> Dict[str, Any]:
     }
 
 
-def _create_conditional_failed_result(i: int, error_msg: str) -> Dict[str, Any]:
+def _create_conditional_failed_result(i: int, error_msg: str, priority_pharm_indices: Optional[list] = None) -> Dict[str, Any]:
     """Create a result dict for failed conditional evaluations with all required fields."""
     base_failed_results = _create_failed_result(i, error_msg)
-    return {
+    res = {
         **base_failed_results,
         # Additional Conditional 3D similarity attributes
         'sim_surf_target': np.nan,
@@ -71,6 +71,12 @@ def _create_conditional_failed_result(i: int, error_msg: str) -> Dict[str, Any]:
         'sim_surf_target_relax_esp_aligned': np.nan,
         'sim_pharm_target_relax_esp_aligned': np.nan,
     }
+    if priority_pharm_indices is not None:
+        res.update({
+            'sim_pharm_priority_target_relax_optimal': np.nan,
+            'sim_pharm_nonpriority_target_relax_optimal': np.nan,
+        })
+    return res
 
 
 def _create_consistency_failed_result(i: int, error_msg: str) -> Dict[str, Any]:
@@ -310,7 +316,8 @@ def _eval_conditional_single(i: int,
                              atoms: np.ndarray,
                              positions: np.ndarray,
                              solvent: Optional[str],
-                             num_processes: int) -> Dict[str, Any]:
+                             num_processes: int,
+                             priority_pharm_indices: Optional[list] = None) -> Dict[str, Any]:
     """
     Evaluate a single molecule and preserve necessary attributes for the pipeline while avoiding
     pickling issues.
@@ -331,8 +338,9 @@ def _eval_conditional_single(i: int,
             condition=condition,
             num_surf_points=num_surf_points,
             pharm_multi_vector=pharm_multi_vector,
+            priority_pharm_indices=priority_pharm_indices,
             num_processes=num_processes,
-            solvent=solvent
+            solvent=solvent,
         )
 
         res = {
@@ -368,12 +376,17 @@ def _eval_conditional_single(i: int,
             'sim_pharm_target_relax_esp_aligned': cond_eval.sim_pharm_target_relax_esp_aligned if cond_eval.sim_pharm_target_relax_esp_aligned is not None else np.nan,
             'error': None
         }
+        if priority_pharm_indices is not None:
+            res.update({
+                'sim_pharm_priority_target_relax_optimal': cond_eval.sim_pharm_priority_target_relax_optimal if cond_eval.sim_pharm_priority_target_relax_optimal is not None else np.nan,
+                'sim_pharm_nonpriority_target_relax_optimal': cond_eval.sim_pharm_nonpriority_target_relax_optimal if cond_eval.sim_pharm_nonpriority_target_relax_optimal is not None else np.nan,
+            })
         return res
 
     except Exception as e:
         error_msg = f"Conditional evaluation failed for molecule {i}: {str(e)}"
         logger.error(f"{error_msg}\n{traceback.format_exc()}")
-        return _create_conditional_failed_result(i, error_msg)
+        return _create_conditional_failed_result(i, error_msg, priority_pharm_indices)
 
 
 def _eval_consistency_single(i: int,

shepherd_score/evaluations/evaluate/evals.py

@@ -473,20 +473,25 @@ def _align_with_esp(self, mp_ref_and_relaxed: MoleculePair) -> float:
 
     def _align_with_pharm(self, mp_ref_and_relaxed: MoleculePair) -> float:
         """
-        Align relaxed molecule to reference/target molecule with pharmacophores
+        Align relaxed molecule to reference/target molecule with pharmacophores.
+
+        Stores aligned fit anchors and vectors on ``self._aligned_pharm_ancs``
+        and ``self._aligned_pharm_vecs`` for downstream subset scoring.
 
         Returns
         -------
         float : Pharmacophore similarity score of optimally aligned molecule.
         """
-        aligned_fit_anchors, aligned_vectors = mp_ref_and_relaxed.align_with_pharm(
+        aligned_fit_anchors, aligned_fit_vectors = mp_ref_and_relaxed.align_with_pharm(
             similarity='tanimoto',
             extended_points=False,
             only_extended=False,
             num_repeats=1,
             trans_init=False,
             use_jax=False
         )
+        self._aligned_pharm_ancs = aligned_fit_anchors
+        self._aligned_pharm_vecs = aligned_fit_vectors
         pharm_similarity = mp_ref_and_relaxed.sim_aligned_pharm
         return float(pharm_similarity)
 
@@ -501,8 +506,10 @@ def __init__(self,
                  condition: str,
                  num_surf_points: int = 400,
                  pharm_multi_vector: Optional[bool] = None,
+                 priority_pharm_indices: Optional[list] = None,
                  solvent: Optional[str] = None,
-                 num_processes: int = 1):
+                 num_processes: int = 1,
+                 ):
         """
         Evaluation pipeline for conditionally-generated molecules.
 
@@ -536,6 +543,15 @@ def __init__(self,
             Number of surface points to sample for similarity scoring. Default is 400.
         pharm_multi_vector : bool, optional
             Use multiple vectors to represent Aro/HBA/HBD or single.
+        priority_pharm_indices : list of int, optional
+            Indices (into ``ref_molec`` pharmacophore arrays) of "priority"
+            pharmacophores. When provided, two additional Tversky
+            (``'tversky_ref'``) scores are computed after the full-set pharm
+            alignment: one for the priority subset and one for the non-priority
+            complement subset, each scored against the full pharmacophore set
+            of the aligned generated molecule. Requires ``condition`` to be
+            ``'pharm'`` or ``'all'`` and ``pharm_multi_vector`` to be a bool.
+            Must satisfy ``0 < len(priority_pharm_indices) < N_pharm``.
         solvent : str, optional
             Solvent type for xTB relaxation.
         num_processes : int, optional
@@ -574,6 +590,11 @@ def __init__(self,
         self.sim_surf_target_relax_esp_aligned = None
         self.sim_pharm_target_relax_esp_aligned = None
 
+        self.sim_pharm_priority_target_relax_optimal = None
+        self.sim_pharm_nonpriority_target_relax_optimal = None
+
+        self.priority_pharm_indices = priority_pharm_indices
+
         # Scoring parameters
         self.num_surf_points = num_surf_points
         self.alpha = ALPHA(self.num_surf_points) # Fitted to probe_radius=1.2
@@ -674,6 +695,39 @@ def __init__(self,
             if (self.condition == 'pharm' or self.condition == 'all') and isinstance(pharm_multi_vector, bool):
                 self.sim_pharm_target_relax_optimal = self._align_with_pharm(mp_ref_and_relaxed=mp_ref_and_relaxed)
 
+                # Priority subset Tversky scoring using the alignment from the full-set pharm alignment above
+                if (self.priority_pharm_indices is not None
+                        and hasattr(self, '_aligned_pharm_ancs')
+                        and self.ref_molec.pharm_ancs is not None):
+                    n_pharm = len(self.ref_molec.pharm_types)
+                    nonpriority_indices = sorted(set(range(n_pharm)) - set(self.priority_pharm_indices))
+                    priority_idx = self.priority_pharm_indices
+
+                    self.sim_pharm_priority_target_relax_optimal = float(get_overlap_pharm_np(
+                        ptype_1=self.ref_molec.pharm_types[priority_idx],
+                        ptype_2=self.molec_post_opt.pharm_types,
+                        anchors_1=self.ref_molec.pharm_ancs[priority_idx],
+                        anchors_2=self._aligned_pharm_ancs,
+                        vectors_1=self.ref_molec.pharm_vecs[priority_idx],
+                        vectors_2=self._aligned_pharm_vecs,
+                        similarity='tversky_ref',
+                        extended_points=False,
+                        only_extended=False
+                    ))
+
+                    if nonpriority_indices:
+                        self.sim_pharm_nonpriority_target_relax_optimal = float(get_overlap_pharm_np(
+                            ptype_1=self.ref_molec.pharm_types[nonpriority_indices],
+                            ptype_2=self.molec_post_opt.pharm_types,
+                            anchors_1=self.ref_molec.pharm_ancs[nonpriority_indices],
+                            anchors_2=self._aligned_pharm_ancs,
+                            vectors_1=self.ref_molec.pharm_vecs[nonpriority_indices],
+                            vectors_2=self._aligned_pharm_vecs,
+                            similarity='tversky_ref',
+                            extended_points=False,
+                            only_extended=False
+                        ))
+
             # Compute ESP-aligned surf and pharmacophore similarity scores
             if mp_ref_and_relaxed.transform_esp is not None and self.condition in ('esp', 'all'):
                 molec_post_opt_esp_aligned = mp_ref_and_relaxed.get_transformed_molecule(mp_ref_and_relaxed.transform_esp)
@@ -733,19 +787,24 @@ def _align_with_esp(self, mp_ref_and_relaxed: MoleculePair) -> float:
 
     def _align_with_pharm(self, mp_ref_and_relaxed: MoleculePair) -> float:
         """
-        Align relaxed molecule to reference/target molecule with pharmacophores
+        Align relaxed molecule to reference/target molecule with pharmacophores.
+
+        Stores aligned fit anchors and vectors on ``self._aligned_pharm_ancs``
+        and ``self._aligned_pharm_vecs`` for downstream subset scoring.
 
         Returns
         -------
         float : Pharmacophore similarity score of optimally aligned molecule.
         """
-        aligned_fit_anchors, aligned_vectors = mp_ref_and_relaxed.align_with_pharm(
+        aligned_fit_anchors, aligned_fit_vectors = mp_ref_and_relaxed.align_with_pharm(
             similarity='tanimoto',
             extended_points=False,
             only_extended=False,
             num_repeats=1,
             trans_init=False,
             use_jax=False
         )
+        self._aligned_pharm_ancs = aligned_fit_anchors
+        self._aligned_pharm_vecs = aligned_fit_vectors
         pharm_similarity = mp_ref_and_relaxed.sim_aligned_pharm
         return float(pharm_similarity)

shepherd_score/evaluations/evaluate/pipelines.py

@@ -387,6 +387,7 @@ def __init__(self,
                  num_surf_points: int = 400,
                  pharm_multi_vector: Optional[bool] = None,
                  solvent: Optional[str] = None,
+                 priority_pharm_indices: Optional[list] = None,
                  ):
         """
         Initialize attributes for conditional evaluation pipeline.
@@ -412,10 +413,21 @@ def __init__(self,
             ref_molec should match.
         solvent : str, optional
             Solvent type for xtb relaxation.
+        priority_pharm_indices : list of int, optional
+            Indices (into ``ref_molec`` pharmacophore arrays) of "priority"
+            pharmacophores. When provided, two additional Tversky
+            (``'tversky_ref'``) scores are computed after the full-set pharm
+            alignment: one for the priority subset
+            (``sims_pharm_priority_target_relax_optimal``) and one for the
+            non-priority complement subset
+            (``sims_pharm_nonpriority_target_relax_optimal``). Requires
+            ``condition`` to be ``'pharm'`` or ``'all'`` and
+            ``pharm_multi_vector`` to be a bool.
         """
         self.generated_mols = generated_mols
         self.num_generated_mols = len(self.generated_mols)
         self.solvent = solvent
+        self.priority_pharm_indices = priority_pharm_indices
 
         self.pharm_multi_vector = pharm_multi_vector
         self.condition = condition
@@ -487,6 +499,10 @@ def __init__(self,
         self.sims_surf_target_relax_esp_aligned = np.empty(self.num_generated_mols)
         self.sims_pharm_target_relax_esp_aligned = np.empty(self.num_generated_mols)
 
+        if self.priority_pharm_indices is not None:
+            self.sims_pharm_priority_target_relax_optimal = np.empty(self.num_generated_mols)
+            self.sims_pharm_nonpriority_target_relax_optimal = np.empty(self.num_generated_mols)
+
         # 2D similarities
         self.graph_similarities = np.empty(self.num_generated_mols)
         self.graph_similarities_post_opt = np.empty(self.num_generated_mols)
@@ -532,7 +548,7 @@ def evaluate(self,
         if num_workers > 1:
             multiprocessing.set_start_method(mp_context, force=True)
             with set_thread_limits(num_processes):
-                inputs = [(i, self.ref_molec, self.condition, self.num_surf_points, self.pharm_multi_vector, atoms, positions, self.solvent, 1)
+                inputs = [(i, self.ref_molec, self.condition, self.num_surf_points, self.pharm_multi_vector, atoms, positions, self.solvent, 1, self.priority_pharm_indices)
                         for i, (atoms, positions) in enumerate(self.generated_mols)]
                 with multiprocessing.Pool(num_workers) as pool:
                     if verbose:
@@ -574,7 +590,8 @@ def evaluate(self,
 
                 res = _eval_conditional_single(
                     i, self.ref_molec, self.condition, self.num_surf_points,
-                    self.pharm_multi_vector, atoms, positions, self.solvent, num_processes
+                    self.pharm_multi_vector, atoms, positions, self.solvent, num_processes,
+                    self.priority_pharm_indices
                 )
                 self._process_single_result(res, i)
 
@@ -658,6 +675,10 @@ def _process_single_result(self, res: Dict[str, Any], i: int):
         self.sims_surf_target_relax_esp_aligned[i] = res['sim_surf_target_relax_esp_aligned']
         self.sims_pharm_target_relax_esp_aligned[i] = res['sim_pharm_target_relax_esp_aligned']
 
+        if self.priority_pharm_indices is not None:
+            self.sims_pharm_priority_target_relax_optimal[i] = res['sim_pharm_priority_target_relax_optimal']
+            self.sims_pharm_nonpriority_target_relax_optimal[i] = res['sim_pharm_nonpriority_target_relax_optimal']
+
 
     def resampling_surf_scores(self) -> Union[np.ndarray, None]:
         """
@@ -767,7 +788,7 @@ def to_pandas(self) -> Tuple[pd.Series, pd.DataFrame]:
         for key, value in self.__dict__.items():
             if key in ('smiles', 'smiles_post_opt', 'morgan_fps', 'morgan_fps_post_opt', 'ref_molec'):
                 continue
-            elif key in ('ref_surf_resampling_scores', 'ref_surf_esp_resampling_scores'):
+            elif key in ('ref_surf_resampling_scores', 'ref_surf_esp_resampling_scores', 'priority_pharm_indices'):
                 global_attrs[key] = value
 
             elif isinstance(value, (list, tuple, np.ndarray)) and not (isinstance(value, np.ndarray) and value.ndim == 0):