import typing as T from abc import ABC, abstractmethod from copy import deepcopy import torch from torchmetrics import Metric, MeanMetric from torchmetrics.functional.retrieval.hit_rate import retrieval_hit_rate from massspecgym.models.base import MassSpecGymModel, Stage from massspecgym.simulation_utils.misc_utils import safelog from massspecgym.simulation_utils.spec_utils import sparse_cosine_distance, \ get_ints_transform_func, get_ints_untransform_func, batched_l1_normalize, \ sparse_jensen_shannon_similarity from massspecgym.simulation_utils.nn_utils import build_lr_scheduler from massspecgym.utils import batch_ptr_to_batch_idx from torch_geometric.utils import unbatch def get_cos_sim_fn(transform_fn, mz_max, mz_bin_res): def cos_sim_fn( true_mzs: torch.Tensor, true_logprobs: torch.Tensor, true_batch_idxs: torch.Tensor, pred_mzs: torch.Tensor, pred_logprobs: torch.Tensor, pred_batch_idxs: torch.Tensor ) -> torch.Tensor: true_logprobs = transform_fn(true_logprobs, true_batch_idxs) pred_logprobs = transform_fn(pred_logprobs, pred_batch_idxs) cos_sims = 1.-sparse_cosine_distance( pred_mzs=pred_mzs, pred_logprobs=pred_logprobs, pred_batch_idxs=pred_batch_idxs, true_mzs=true_mzs, true_logprobs=true_logprobs, true_batch_idxs=true_batch_idxs, mz_max=mz_max, mz_bin_res=mz_bin_res ) return cos_sims return deepcopy(cos_sim_fn) def get_js_sim_fn(transform_fn, mz_max, mz_bin_res): def js_sim_fn( true_mzs: torch.Tensor, true_logprobs: torch.Tensor, true_batch_idxs: torch.Tensor, pred_mzs: torch.Tensor, pred_logprobs: torch.Tensor, pred_batch_idxs: torch.Tensor ) -> torch.Tensor: true_logprobs = transform_fn(true_logprobs, true_batch_idxs) pred_logprobs = transform_fn(pred_logprobs, pred_batch_idxs) js_sims = sparse_jensen_shannon_similarity( pred_mzs=pred_mzs, pred_logprobs=pred_logprobs, pred_batch_idxs=pred_batch_idxs, true_mzs=true_mzs, true_logprobs=true_logprobs, true_batch_idxs=true_batch_idxs, mz_max=mz_max, mz_bin_res=mz_bin_res ) return js_sims return deepcopy(js_sim_fn) class SimulationMassSpecGymModel(MassSpecGymModel, ABC): def __init__(self, **kwargs): super().__init__(**kwargs) # call setup functions self._setup_model() self._setup_loss_fn() self._setup_spec_fns() self._setup_metric_fns() @abstractmethod def _setup_model(self): pass def configure_optimizers(self): if self.hparams.optimizer_type == "adam": optimizer_cls = torch.optim.Adam elif self.hparams.optimizer_type == "adamw": optimizer_cls = torch.optim.AdamW elif self.hparams.optimizer_type == "sgd": optimizer_cls = torch.optim.SGD else: raise ValueError(f"Unknown optimizer {self.optimizer}") optimizer = optimizer_cls( self.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay ) ret = { "optimizer": optimizer, } if self.hparams.lr_schedule: scheduler = build_lr_scheduler( optimizer=optimizer, decay_rate=self.hparams.lr_decay_rate, warmup_steps=self.hparams.lr_warmup_steps, decay_steps=self.hparams.lr_decay_steps, ) ret["lr_scheduler"] = { "scheduler": scheduler, "frequency": 1, "interval": "step", } return ret def _setup_spec_fns(self): self.ints_transform_func = get_ints_transform_func(self.hparams.ints_transform) self.ints_untransform_func = get_ints_untransform_func(self.hparams.ints_transform) self.ints_normalize_func = batched_l1_normalize def _preproc_spec(self,spec_mzs,spec_ints,spec_batch_idxs): # transform spec_ints = spec_ints * 1000. spec_ints = self.ints_transform_func(spec_ints) # renormalize spec_ints = self.ints_normalize_func( spec_ints, spec_batch_idxs ) spec_ints = safelog(spec_ints) return spec_mzs, spec_ints, spec_batch_idxs def _setup_loss_fn(self): def _loss_fn( true_mzs: torch.Tensor, true_logprobs: torch.Tensor, true_batch_idxs: torch.Tensor, pred_mzs: torch.Tensor, pred_logprobs: torch.Tensor, pred_batch_idxs: torch.Tensor ): cos_dist = sparse_cosine_distance( true_mzs=true_mzs, true_logprobs=true_logprobs, true_batch_idxs=true_batch_idxs, pred_mzs=pred_mzs, pred_logprobs=pred_logprobs, pred_batch_idxs=pred_batch_idxs, mz_max=self.hparams.mz_max, mz_bin_res=self.hparams.mz_bin_res ) return cos_dist self.loss_fn = _loss_fn def _setup_metric_fns(self): # untransformed def no_transform_fn(logprobs, batch_idxs): probs = self.ints_untransform_func(torch.exp(logprobs), batch_idxs) logprobs = safelog(self.ints_normalize_func(probs, batch_idxs)) return logprobs self.cos_sim_fn = get_cos_sim_fn( deepcopy(no_transform_fn), self.hparams.mz_max, self.hparams.mz_bin_res ) self.js_sim_fn = get_js_sim_fn( deepcopy(no_transform_fn), self.hparams.mz_max, self.hparams.mz_bin_res ) # sqrt transform def sqrt_transform_fn(logprobs, batch_idxs): probs = self.ints_untransform_func(torch.exp(logprobs), batch_idxs) probs = probs**0.5 logprobs = safelog(self.ints_normalize_func(probs, batch_idxs)) return logprobs self.cos_sim_sqrt_fn = get_cos_sim_fn( deepcopy(sqrt_transform_fn), self.hparams.mz_max, self.hparams.mz_bin_res ) # obj transform obj_transform_fn = lambda x, y: x self.cos_sim_obj_fn = get_cos_sim_fn( deepcopy(obj_transform_fn), self.hparams.mz_max, self.hparams.mz_bin_res ) def forward(self, **kwargs) -> dict: return self.model.forward(**kwargs) def get_retrieval_outputs( self, true_mzs, true_logprobs, true_batch_idxs, candidates_data, **kwargs ) -> dict: input_data = {} for k in candidates_data: input_data[k] = candidates_data[k] c_ptr = input_data["batch_ptr"] for k in ["precursor_mz", "adduct", "instrument_type", "collision_energy"]: if k in kwargs: assert not k in input_data, k input_data[k] = torch.repeat_interleave(kwargs[k],c_ptr) output_data = self.forward(**input_data) c_offset = torch.cumsum(c_ptr, dim=0) c_offset = torch.cat([ torch.zeros((1,),dtype=c_offset.dtype,device=c_offset.device), c_offset ], dim=0) c_total = torch.sum(c_ptr) c_true_mzs, c_true_logprobs, c_true_batch_idxs = [], [], [] # TODO: replace with unbatch? for idx in range(c_ptr.shape[0]): mask = true_batch_idxs==idx ptr = c_ptr[idx] c_true_mzs.append(true_mzs[mask].repeat(ptr)) c_true_logprobs.append(true_logprobs[mask].repeat(ptr)) c_true_batch_idxs.append(torch.repeat_interleave(torch.arange(ptr,device=c_ptr.device),mask.sum()) + c_offset[idx]) c_true_mzs = torch.cat(c_true_mzs, dim=0) c_true_logprobs = torch.cat(c_true_logprobs, dim=0) c_true_batch_idxs = torch.cat(c_true_batch_idxs, dim=0) assert output_data["pred_batch_idxs"].max() == c_true_batch_idxs.max() c_scores = self.cos_sim_fn( true_mzs=c_true_mzs, true_logprobs=c_true_logprobs, true_batch_idxs=c_true_batch_idxs, pred_mzs=output_data["pred_mzs"], pred_logprobs=output_data["pred_logprobs"], pred_batch_idxs=output_data["pred_batch_idxs"] ) c_labels = input_data["labels"] assert c_scores.shape[0] == c_labels.shape[0] == c_total, (c_scores.shape, c_labels.shape, c_total) # prepare dictionary ret_out_d = {} ret_out_d["retrieval_scores"] = c_scores ret_out_d["retrieval_labels"] = c_labels ret_out_d["retrieval_batch_ptr"] = c_ptr return ret_out_d def step( self, batch: dict, stage: Stage = Stage.NONE ) -> dict: true_mzs, true_logprobs, true_batch_idxs = self._preproc_spec( batch["spec_mzs"], batch["spec_ints"], batch["spec_batch_idxs"] ) out_d = self.forward( **batch ) pred_mzs = out_d["pred_mzs"] pred_logprobs = out_d["pred_logprobs"] pred_batch_idxs = out_d["pred_batch_idxs"] loss = self.loss_fn( true_mzs=true_mzs, true_logprobs=true_logprobs, true_batch_idxs=true_batch_idxs, pred_mzs=pred_mzs, pred_logprobs=pred_logprobs, pred_batch_idxs=pred_batch_idxs ) mean_loss = torch.mean(loss) out_d = { "loss": mean_loss, "losses": loss, "pred_mzs": pred_mzs, "pred_logprobs": pred_logprobs, "pred_batch_idxs": pred_batch_idxs, "true_mzs": true_mzs, "true_logprobs": true_logprobs, "true_batch_idxs": true_batch_idxs } if stage == Stage.TEST and self.hparams.do_retrieval: # retrieval assert "candidates_data" in batch, batch.keys() ret_out_d = self.get_retrieval_outputs( true_mzs=true_mzs, true_logprobs=true_logprobs, true_batch_idxs=true_batch_idxs, **batch ) out_d.update(ret_out_d) return out_d def on_batch_end( self, outputs: T.Any, batch: dict, batch_idx: int, stage: Stage ) -> None: """ Compute evaluation metrics for the retrieval model based on the batch and corresponding predictions. This method will be used in the `on_train_batch_end`, `on_validation_batch_end`, since `on_test_batch_end` is overridden below. """ batch_size = torch.max(outputs["true_batch_idxs"])+1 metric_vals = {} metric_vals["losses"] = outputs["losses"] # Log loss self.log( stage.to_pref() + "loss_step", outputs["loss"], batch_size=batch_size, sync_dist=True, prog_bar=True, on_step=True, on_epoch=False ) if stage not in self.log_only_loss_at_stages: metric_vals.update(self.evaluate_similarity_step( pred_mzs=outputs["pred_mzs"], pred_logprobs=outputs["pred_logprobs"], pred_batch_idxs=outputs["pred_batch_idxs"], true_mzs=outputs["true_mzs"], true_logprobs=outputs["true_logprobs"], true_batch_idxs=outputs["true_batch_idxs"], stage=stage )) if stage == stage.TEST and self.df_test_path is not None: self._update_df_test(metric_vals) def on_test_batch_end( self, outputs: T.Any, batch: dict, batch_idx: int ) -> None: stage = Stage.TEST self.on_batch_end( outputs=outputs, batch=batch, batch_idx=batch_idx, stage=stage ) if stage not in self.log_only_loss_at_stages and self.hparams.do_retrieval: metric_vals = {} metric_vals |= self.evaluate_retrieval_step( scores=outputs["retrieval_scores"], labels=outputs["retrieval_labels"], batch_ptr=outputs["retrieval_batch_ptr"], stage=stage ) if self.df_test_path is not None: self._update_df_test(metric_vals) def evaluate_similarity_step( self, pred_mzs: torch.Tensor, pred_logprobs: torch.Tensor, pred_batch_idxs: torch.Tensor, true_mzs: torch.Tensor, true_logprobs: torch.Tensor, true_batch_idxs: torch.Tensor, stage: Stage ) -> dict: metric_vals = {} batch_size = torch.max(true_batch_idxs).item()+1 assert "cos_sim" in self.hparams.sim_metrics, self.hparams.sim_metrics for metric_name in self.hparams.sim_metrics: metric_fn = getattr(self, metric_name+"_fn") metric = metric_fn( true_mzs=true_mzs, true_logprobs=true_logprobs, true_batch_idxs=true_batch_idxs, pred_mzs=pred_mzs, pred_logprobs=pred_logprobs, pred_batch_idxs=pred_batch_idxs ) self._update_metric( name=stage.to_pref() + metric_name, metric_class=MeanMetric, update_args=(metric,), batch_size=batch_size, prog_bar=True, log=True, log_n_samples=False, bootstrap=(stage == Stage.TEST) ) metric_vals[metric_name] = metric return metric_vals def evaluate_retrieval_step( self, scores: torch.Tensor, labels: torch.Tensor, batch_ptr: torch.Tensor, stage: Stage, ) -> dict: """ Main evaluation method for the retrieval models. The retrieval step is evaluated by computing the hit rate at different top-k values. Args: scores (torch.Tensor): Concatenated scores for all candidates for all samples in the batch labels (torch.Tensor): Concatenated True/False labels for all candidates for all samples in the batch batch_ptr (torch.Tensor): Number of each sample's candidates in the concatenated tensors """ assert stage == Stage.TEST, stage # Initialize return dictionary to store metric values per sample metric_vals = {} # Evaluate hitrate at different top-k values indexes = batch_ptr_to_batch_idx(batch_ptr) scores = scores + 1e-5*torch.randn_like(scores) scores = unbatch(scores, indexes) labels = unbatch(labels, indexes) for at_k in self.hparams.at_ks: hit_rates = [] for scores_sample, labels_sample in zip(scores, labels): hit_rates.append(retrieval_hit_rate(scores_sample, labels_sample, top_k=at_k)) hit_rates = torch.tensor(hit_rates, device=batch_ptr.device) metric_name = f"{stage.to_pref()}hit_rate@{at_k}" self._update_metric( metric_name, MeanMetric, (hit_rates,), batch_size=batch_ptr.size(0), bootstrap=(stage == Stage.TEST) ) metric_vals[metric_name] = hit_rates return metric_vals def on_test_epoch_end(self): self._save_df_test()