import typing as T from abc import ABC import torch import pandas as pd from rdkit import Chem from rdkit.DataStructs import TanimotoSimilarity from torchmetrics.aggregation import MeanMetric from massspecgym.models.base import MassSpecGymModel, Stage from massspecgym.utils import morgan_fp, mol_to_inchi_key, MyopicMCES class DeNovoMassSpecGymModel(MassSpecGymModel, ABC): def __init__( self, top_ks: T.Iterable[int] = (1, 10), myopic_mces_kwargs: T.Optional[T.Mapping] = None, *args, **kwargs ): super().__init__(*args, **kwargs) self.top_ks = top_ks self.myopic_mces = MyopicMCES(**(myopic_mces_kwargs or {})) self.mol_pred_kind: T.Literal["smiles", "rdkit"] = "smiles" # caches of already computed results to avoid expensive re-computations self.mces_cache = dict() self.mol_2_morgan_fp = dict() def on_batch_end( self, outputs: T.Any, batch: dict, batch_idx: int, stage: Stage ) -> None: self.log( f"{stage.to_pref()}loss", outputs['loss'], batch_size=batch['spec'].size(0), sync_dist=True, prog_bar=True, ) if stage in self.log_only_loss_at_stages: return metric_vals = self.evaluate_de_novo_step( outputs["mols_pred"], # (bs, k) list of generated rdkit molecules or SMILES strings batch["mol"], # (bs) list of ground truth SMILES strings stage=stage ) if stage == Stage.TEST and self.df_test_path is not None: self._update_df_test(metric_vals) def evaluate_de_novo_step( self, mols_pred: list[list[T.Optional[Chem.Mol | str]]], mol_true: list[str], stage: Stage, ) -> dict[str, torch.Tensor]: """ # TODO: refactor to compute only for max(k) and then use the result to obtain the rest by subsetting. Main evaluation method for the models for de novo molecule generation from mass spectra. Args: mols_pred (list[list[Mol | str]]): (bs, k) list of generated rdkit molecules or SMILES strings with possible Nones if no molecule was generated mol_true (list[str]): (bs) list of ground-truth SMILES strings """ # Initialize return dictionary to store metric values per sample metric_vals = {} # Get SMILES and RDKit molecule objects for all predictions if self.mol_pred_kind == "smiles": smiles_pred_valid, mols_pred_valid = [], [] for mols_pred_sample in mols_pred: smiles_pred_valid_sample, mols_pred_valid_sample = [], [] for s in mols_pred_sample: m = Chem.MolFromSmiles(s) if s is not None else None # If SMILES cannot be converted to RDKit molecule, the molecule is set to None smiles_pred_valid_sample.append(s if m is not None else None) mols_pred_valid_sample.append(m) smiles_pred_valid.append(smiles_pred_valid_sample) mols_pred_valid.append(mols_pred_valid_sample) smiles_pred, mols_pred = smiles_pred_valid, mols_pred_valid elif self.mol_pred_kind == "rdkit": smiles_pred = [ [Chem.MolToSmiles(m) if m is not None else None for m in ms] for ms in mols_pred ] else: raise ValueError(f"Invalid mol_pred_kind: {self.mol_pred_kind}") # Auxiliary metric: number of valid molecules self._update_metric( stage.to_pref() + f"num_valid_mols", MeanMetric, ([sum([m is not None for m in ms]) for ms in mols_pred],), batch_size=len(mols_pred), ) # Get RDKit molecule objects for ground truth smile_true = mol_true mol_true = [Chem.MolFromSmiles(sm) for sm in mol_true] def _get_morgan_fp_with_cache(mol): """ A helper function to retrieve either cached Morgan Fingerprint value, or to compute and cache it @param mol: RDKit molecule object @return: """ if mol not in self.mol_2_morgan_fp: self.mol_2_morgan_fp[mol] = morgan_fp(mol, to_np=False) return self.mol_2_morgan_fp[mol] # Evaluate top-k metrics for top_k in self.top_ks: # Get top-k predicted molecules for each ground-truth sample smiles_pred_top_k = [smiles_pred_sample[:top_k] for smiles_pred_sample in smiles_pred] mols_pred_top_k = [mols_pred_sample[:top_k] for mols_pred_sample in mols_pred] # 1. Evaluate minimum common edge subgraph: # Calculate MCES distance between top-k predicted molecules and ground truth and # report the minimum distance. The minimum distances for each sample in the batch are # averaged across the epoch. if stage not in self.no_mces_metrics_at_stages: min_mces_dists = [] mces_thld = 100 # Iterate over batch for preds, true in zip(smiles_pred_top_k, smile_true): # Iterate over top-k predicted molecule samples dists = [] for pred in preds: if pred is None: dists.append(mces_thld) else: if (true, pred) not in self.mces_cache: mce_val = self.myopic_mces(true, pred) self.mces_cache[(true, pred)] = mce_val dists.append(self.mces_cache[(true, pred)]) min_mces_dists.append(min(min(dists), mces_thld)) min_mces_dists = torch.tensor(min_mces_dists, device=self.device) # Log metric_name = stage.to_pref() + f"top_{top_k}_mces_dist" self._update_metric( metric_name, MeanMetric, (min_mces_dists,), batch_size=len(min_mces_dists), bootstrap=stage == Stage.TEST ) metric_vals[metric_name] = min_mces_dists # 2. Evaluate Tanimoto similarity: # Calculate Tanimoto similarity between top-k predicted molecules and ground truth and # report the maximum similarity. The maximum similarities for each sample in the batch # are averaged across the epoch. fps_pred_top_k = [ [_get_morgan_fp_with_cache(m) if m is not None else None for m in ms] for ms in mols_pred_top_k ] fp_true = [_get_morgan_fp_with_cache(m) for m in mol_true] max_tanimoto_sims = [] # Iterate over batch for preds, true in zip(fps_pred_top_k, fp_true): # Iterate over top-k predicted molecule samples sims = [ TanimotoSimilarity(true, pred) if pred is not None else 0 for pred in preds ] max_tanimoto_sims.append(max(sims)) max_tanimoto_sims = torch.tensor(max_tanimoto_sims, device=self.device) # Log metric_name = stage.to_pref() + f"top_{top_k}_max_tanimoto_sim" self._update_metric( metric_name, MeanMetric, (max_tanimoto_sims,), batch_size=len(max_tanimoto_sims), bootstrap=stage == Stage.TEST ) metric_vals[metric_name] = max_tanimoto_sims # 3. Evaluate exact match (accuracy): # Calculate if the ground truth molecule is in the top-k predicted molecules and report # the average across the epoch. in_top_k = [ mol_to_inchi_key(true) in [ mol_to_inchi_key(pred) if pred is not None else None for pred in preds ] for true, preds in zip(mol_true, mols_pred_top_k) ] in_top_k = torch.tensor(in_top_k, device=self.device) # Log metric_name = stage.to_pref() + f"top_{top_k}_accuracy" self._update_metric( metric_name, MeanMetric, (in_top_k,), batch_size=len(in_top_k), bootstrap=stage == Stage.TEST ) metric_vals[metric_name] = in_top_k return metric_vals def test_step( self, batch: dict, batch_idx: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: outputs = super().test_step(batch, batch_idx) # Get generated (i.e., predicted) SMILES if self.df_test_path is not None: self._update_df_test({ 'identifier': batch['identifier'], 'mols_pred': outputs['mols_pred'] }) return outputs def on_test_epoch_end(self): # Save test data frame to disk if self.df_test_path is not None: df_test = pd.DataFrame(self.df_test) self.df_test_path.parent.mkdir(parents=True, exist_ok=True) df_test.to_pickle(self.df_test_path)