Spaces:
Sleeping
Sleeping
| import pandas as pd | |
| import json | |
| import typing as T | |
| import numpy as np | |
| import torch | |
| import massspecgym.utils as utils | |
| from pathlib import Path | |
| from torch.utils.data.dataset import Dataset | |
| from torch.utils.data.dataloader import default_collate | |
| import dgl | |
| from collections import defaultdict | |
| from massspecgym.data.transforms import SpecTransform, MolTransform, MolToInChIKey | |
| from massspecgym.data.datasets import MassSpecDataset | |
| import flare.utils.data as data_utils | |
| from torch.nn.utils.rnn import pad_sequence | |
| from massspecgym.models.base import Stage | |
| import pickle | |
| import math | |
| import itertools | |
| from rdkit.Chem import AllChem | |
| from rdkit import Chem | |
| from magma.run_magma import run_magma | |
| import matchms | |
| class JESTR1_MassSpecDataset(MassSpecDataset): | |
| def __init__( | |
| self, | |
| spectra_view: str, | |
| fp_dir_pth: str = None, | |
| cons_spec_dir_pth: str = None, | |
| NL_spec_dir_pth: str = None, | |
| **kwargs | |
| ): | |
| super().__init__(**kwargs) | |
| self.use_fp = False | |
| self.use_cons_spec = False | |
| self.use_NL_spec = False | |
| self.spectra_view = spectra_view | |
| # load fingerprints | |
| self._load_fp(fp_dir_pth) | |
| # load consensus | |
| self._load_cons_spec(cons_spec_dir_pth) | |
| # load NL specs | |
| self._load_NL_spec(NL_spec_dir_pth) | |
| def _load_fp(self, fp_dir_pth): | |
| if fp_dir_pth is not None: | |
| self.use_fp = True | |
| if fp_dir_pth: | |
| with open(fp_dir_pth, 'rb') as f: | |
| self.smiles_to_fp = pickle.load(f) | |
| else: | |
| self.smiles_to_fp = {} | |
| def _load_cons_spec(self, cons_spec_dir_pth): | |
| if cons_spec_dir_pth is not None: | |
| self.use_cons_spec = True | |
| with open(cons_spec_dir_pth, 'rb') as f: | |
| cons_specs = pickle.load(f) | |
| # Convert spectra to matchms spectra | |
| matchMS_preparer = data_utils.PrepMatchMS(self.spectra_view) | |
| spectra = cons_specs.apply(matchMS_preparer.prepare,axis=1) | |
| self.cons_specs = dict(zip(cons_specs['smiles'].tolist(), spectra)) | |
| def _load_NL_spec(self, NL_spec_dir_pth): | |
| if NL_spec_dir_pth is not None: | |
| self.use_NL_spec = True | |
| with open(NL_spec_dir_pth, 'rb') as f: | |
| NL_specs = pickle.load(f) | |
| # Convert spectra to matchms spectra | |
| matchMS_preparer = data_utils.PrepMatchMS(self.spectra_view) | |
| self.NL_specs = NL_specs.apply(matchMS_preparer.prepare,axis=1) | |
| def __getitem__(self, i, transform_spec: bool = True, transform_mol: bool = True): | |
| spec = self.spectra[i] | |
| metadata = self.metadata.iloc[i] | |
| mol = metadata["smiles"] if 'smiles' in metadata else metadata["identifier"] | |
| # Apply all transformations to the spectrum | |
| item = {} | |
| if transform_spec and self.spec_transform: | |
| if isinstance(self.spec_transform, dict): | |
| for key, transform in self.spec_transform.items(): | |
| item[key] = transform(spec) if transform is not None else spec | |
| else: | |
| item["spec"] = self.spec_transform(spec) | |
| if self.return_mol_freq: | |
| item["mol_freq"] = metadata["mol_freq"] | |
| if self.return_identifier: | |
| item["identifier"] = metadata["identifier"] | |
| if self.use_fp and self.smiles_to_fp: | |
| item['fp'] = torch.Tensor(self.smiles_to_fp[mol].ToList()) | |
| if self.use_cons_spec: | |
| item['cons_spec'] = self.spec_transform[self.spectra_view](self.cons_specs[mol]) | |
| if self.use_NL_spec: | |
| item['NL_spec'] = self.spec_transform[self.spectra_view](self.NL_specs[i]) | |
| # Apply all transformations to the molecule | |
| if transform_mol and self.mol_transform: | |
| if isinstance(self.mol_transform, dict): | |
| for key, transform in self.mol_transform.items(): | |
| item[key] = transform(mol) if transform is not None else mol | |
| else: | |
| item["mol"] = self.mol_transform(mol) | |
| else: | |
| item["mol"] = mol | |
| return item | |
| class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset): | |
| def __init__( | |
| self, | |
| spectra_view: str, | |
| spec_transform: T.Optional[T.Union[SpecTransform, T.Dict[str, SpecTransform]]], | |
| mol_transform: T.Optional[T.Union[MolTransform, T.Dict[str, MolTransform]]], | |
| pth: T.Optional[Path], | |
| subformula_dir_pth: str, | |
| fp_dir_pth: str = None, | |
| NL_spec_dir_pth: str = None, | |
| cons_spec_dir_pth: str = None, | |
| return_mol_freq: bool = False, | |
| return_identifier: bool = True, | |
| dtype: T.Type = torch.float32, | |
| formula_source = 'default', | |
| stage: Stage = Stage.TRAIN | |
| ): | |
| """ | |
| Args: | |
| """ | |
| self.pth = pth | |
| self.spec_transform = spec_transform | |
| self.mol_transform = mol_transform | |
| self.return_mol_freq = return_mol_freq | |
| self.pred_fp = False | |
| self.use_fp = False | |
| self.use_cons_spec = False | |
| self.use_NL_spec = False | |
| self.spectra_view = spectra_view | |
| self.formula_source = formula_source | |
| self.subformula_dir_pth = subformula_dir_pth | |
| if isinstance(self.pth, str): | |
| self.pth = Path(self.pth) | |
| self.spectra_view = spectra_view | |
| print("Data path: ", self.pth) | |
| self.metadata = pd.read_csv(self.pth, sep="\t") | |
| # load subformulas | |
| id_to_spec = self._load_id_to_spec(stage) | |
| # load fingerprints | |
| self._load_fp(fp_dir_pth) | |
| # load consensus spectra | |
| self._load_cons_spec(cons_spec_dir_pth) | |
| # load NL specs | |
| self._load_NL_spec(NL_spec_dir_pth) | |
| self.metadata = self.metadata[self.metadata['identifier'].isin(id_to_spec)] | |
| formula_df = pd.DataFrame.from_dict(id_to_spec, orient='index').reset_index().rename(columns={'index': 'identifier'}) | |
| self.metadata = self.metadata.merge(formula_df, on='identifier') | |
| # create matchms spectra | |
| matchMS_preparer = data_utils.PrepMatchMS(spectra_view=spectra_view) | |
| self.spectra = self.metadata.apply(matchMS_preparer.prepare,axis=1) | |
| if self.return_mol_freq: | |
| if "inchikey" not in self.metadata.columns: | |
| self.metadata["inchikey"] = self.metadata["smiles"].apply(utils.smiles_to_inchi_key) | |
| self.metadata["mol_freq"] = self.metadata.groupby("inchikey")["inchikey"].transform("count") | |
| self.return_identifier = return_identifier | |
| self.dtype = dtype | |
| def __getitem__(self, i, transform_spec: bool = True, transform_mol: bool = True): | |
| item = super().__getitem__(i, transform_spec, transform_mol = False) | |
| mol = item['mol'] #smiles | |
| # transform mol | |
| if transform_mol: | |
| if isinstance(self.mol_transform, dict): | |
| for key, transform in self.mol_transform.items(): | |
| item[key] = transform(mol) if transform is not None else mol | |
| else: | |
| item["mol"] = self.mol_transform(mol) | |
| return item | |
| def _load_id_to_spec(self, stage): | |
| # if stage == Stage.TRAIN: | |
| # self.metadata = self.metadata[self.metadata['fold'] != Stage.TEST.value] | |
| # else: | |
| # self.metadata = self.metadata[self.metadata['fold'] == Stage.TEST.value] | |
| all_spec_ids = self.metadata['identifier'].tolist() | |
| self.subformulaLoader = data_utils.Subformula_Loader(spectra_view=self.spectra_view, dir_path=self.subformula_dir_pth, formula_source=self.formula_source) | |
| form_list = self.metadata['formula'].tolist() | |
| prec_mz_list = self.metadata['precursor_mz'].tolist() | |
| id_to_spec = self.subformulaLoader(all_spec_ids, form_list, prec_mz_list) | |
| # create subformula spectra if no subformula is available | |
| tmp_ids = [spec_id for spec_id in all_spec_ids if spec_id not in id_to_spec] | |
| tmp_df = self.metadata[self.metadata['identifier'].isin(tmp_ids)] | |
| tmp_df['spec'] = tmp_df.apply(lambda row: data_utils.make_tmp_subformula_spectra(row), axis=1) | |
| id_to_spec.update(dict(zip(tmp_df['identifier'].tolist(), tmp_df['spec'].tolist()))) | |
| return id_to_spec | |
| class ContrastiveDataset(Dataset): | |
| def __init__( | |
| self, | |
| spec_mol_data, | |
| ): | |
| super().__init__() | |
| indices = spec_mol_data.indices | |
| self.spec_mol_data = spec_mol_data | |
| self.smiles_to_specmol_ids = spec_mol_data.dataset.metadata.loc[indices].groupby('smiles').indices | |
| self.smiles_to_spec_couter = defaultdict(int) | |
| self.smiles_list = list(self.smiles_to_specmol_ids.keys()) | |
| def __len__(self) -> int: | |
| return len(self.smiles_list) | |
| def __getitem__(self, i:int) -> dict: | |
| mol = self.smiles_list[i] | |
| # select spectrum (iterate through list of spectra) | |
| specmol_ids = self.smiles_to_specmol_ids[mol] | |
| counter = self.smiles_to_spec_couter[mol] | |
| specmol_id = specmol_ids[counter % len(specmol_ids)] | |
| item = self.spec_mol_data.__getitem__(specmol_id) | |
| self.smiles_to_spec_couter[mol] = counter+1 | |
| # item['smiles'] = mol | |
| # item['spec_id'] = specmol_id | |
| return item | |
| def collate_fn(batch: T.Iterable[dict], spec_enc: str, spectra_view: str, stage=None, batch_mol: bool = True) -> dict: | |
| mol_key = 'cand' if stage == Stage.TEST else 'mol' | |
| non_standard_collate = ['mol', 'cand', 'aug_cands', 'cons_spec', 'aug_cands_fp', 'NL_spec'] | |
| require_pad = False | |
| if 'Formula' in spectra_view or 'Tokens' in spectra_view: | |
| require_pad = True | |
| padding_value=-5 if spec_enc in ('Transformer_Formula', 'Formula_BinnedSpec', 'Transformer_MzInt') else 0 | |
| non_standard_collate.append(spectra_view) | |
| else: | |
| non_standard_collate.remove('cons_spec') | |
| non_standard_collate.remove('NL_spec') | |
| collated_batch = {} | |
| # standard collate | |
| for k in batch[0].keys(): | |
| if k not in non_standard_collate: | |
| try: | |
| collated_batch[k] = default_collate([item[k] for item in batch]) | |
| except: | |
| print(f"Error in collating key {k}") | |
| raise | |
| # batch graphs | |
| if batch_mol: | |
| batch_mol = [] | |
| batch_mol_nodes= [] | |
| for item in batch: | |
| batch_mol.append(item[mol_key]) | |
| batch_mol_nodes.append(item[mol_key].num_nodes()) | |
| collated_batch[mol_key] = dgl.batch(batch_mol) | |
| collated_batch['mol_n_nodes'] = batch_mol_nodes | |
| # pad peaks/formulas | |
| if require_pad: | |
| peaks = [] | |
| n_peaks = [] | |
| for item in batch: | |
| peaks.append(item[spectra_view]) | |
| n_peaks.append(len(item[spectra_view])) | |
| collated_batch[spectra_view] = pad_sequence(peaks, batch_first=True, padding_value=padding_value) | |
| collated_batch['n_peaks'] = n_peaks | |
| if 'cons_spec' in batch[0]: | |
| peaks = [] | |
| n_peaks = [] | |
| for item in batch: | |
| peaks.append(item['cons_spec']) | |
| n_peaks.append(len(item['cons_spec'])) | |
| collated_batch['cons_spec'] = pad_sequence(peaks, batch_first=True, padding_value=padding_value) | |
| collated_batch['cons_n_peaks'] = n_peaks | |
| if 'NL_spec' in batch[0]: | |
| peaks = [] | |
| n_peaks = [] | |
| for item in batch: | |
| peaks.append(item['NL_spec']) | |
| n_peaks.append(len(item['NL_spec'])) | |
| collated_batch['NL_spec'] = pad_sequence(peaks, batch_first=True, padding_value=padding_value) | |
| collated_batch['NL_n_peaks'] = n_peaks | |
| return collated_batch | |
| class ExpandedRetrievalDataset: | |
| '''Used for testing only | |
| Assumes 'fold' column defines the split''' | |
| def __init__(self, | |
| use_formulas: bool = True, | |
| mol_label_transform: MolTransform = MolToInChIKey(), | |
| candidates_pth: T.Optional[T.Union[Path, str]] = None, | |
| fp_size: int = None, | |
| fp_radius: int = None, | |
| use_magma = False, | |
| **kwargs): | |
| self.use_magma = use_magma | |
| self.instance = MassSpecDataset_PeakFormulas(**kwargs, return_mol_freq=False, stage = Stage.TEST) if use_formulas else JESTR1_MassSpecDataset(**kwargs, return_mol_freq=False) | |
| if self.use_fp: | |
| self.fpgen = AllChem.GetMorganGenerator(radius=fp_radius,fpSize=fp_size) | |
| self.candidates_pth = candidates_pth | |
| self.mol_label_transform = mol_label_transform | |
| # Read candidates_pth from json to dict: SMILES -> respective candidate SMILES | |
| with open(self.candidates_pth, "r") as file: | |
| candidates = json.load(file) | |
| self.candidates = {} | |
| for s, cand in candidates.items(): | |
| clean_cands = [] | |
| for c in cand: | |
| try: | |
| if '.' not in c: | |
| clean_cands.append(c) | |
| except: | |
| print(f"Error in processing candidate {c} for smiles {s}") | |
| pass | |
| self.candidates[s] = clean_cands | |
| self.spec_cand = [] #(spec index, cand_smiles, true_label) | |
| # use for external dataset where target smiles is not known | |
| # self.candidates should be a dict of identifier to candidates | |
| if 'smiles' not in self.metadata.columns: | |
| if not isinstance(self.metadata.iloc[0]['identifier'], str): | |
| self.metadata['smiles'] = self.metadata['identifier'].apply(str) | |
| else: | |
| self.metadata['smiles'] = self.metadata['identifier'] | |
| # keep datapoints where there are candidates | |
| self.metadata = self.metadata[self.metadata['smiles'].isin(self.candidates.keys())] | |
| test_smiles = self.metadata[self.metadata['fold'] == "test"]['smiles'].tolist() | |
| test_ms_id = self.metadata[self.metadata['fold'] == "test"]['identifier'].tolist() | |
| self.spec_id_to_index = dict(zip(self.metadata['identifier'], self.metadata.index)) | |
| for spec_id, s in zip(test_ms_id, test_smiles): | |
| candidates = self.candidates[s] | |
| # mol_label = self.mol_label_transform(s) | |
| # labels = [self.mol_label_transform(c) == mol_label for c in candidates] | |
| labels = [c == s for c in candidates] | |
| if len(candidates) == 0: | |
| print(f"Skipping {spec_id}; empty candidate set") | |
| continue | |
| if not any(labels): | |
| # print(f"Target smiles not in candidate set") | |
| pass | |
| self.spec_cand.extend([(self.spec_id_to_index[spec_id], candidates[j], k) for j, k in enumerate(labels)]) | |
| def __getattr__(self, name): | |
| return self.instance.__getattribute__(name) | |
| def __len__(self): | |
| return len(self.spec_cand) | |
| def __getitem__(self, i): | |
| spec_i = self.spec_cand[i][0] | |
| cand_smiles = self.spec_cand[i][1] | |
| label = self.spec_cand[i][2] | |
| if self.use_magma: | |
| item = self.instance.__getitem__(spec_i, transform_mol=False, transform_spec=False) | |
| mzs = np.array([float(x) for x in self.metadata.iloc[spec_i]['mzs'].split(',')]) | |
| intensities = np.array([float(x) for x in self.metadata.iloc[spec_i]['intensities'].split(',')]) | |
| adduct = self.metadata.iloc[spec_i]['adduct'] | |
| precursor_mz = self.metadata.iloc[spec_i]['precursor_mz'] | |
| formula = self.metadata.iloc[spec_i]['formula'] | |
| spec_data = run_magma(i, mzs, intensities, cand_smiles, adduct) | |
| spec = self.subformulaLoader.load_magma_data(spec_data, formula, precursor_mz) | |
| spec = matchms.Spectrum( | |
| mz = np.array(spec['formula_mzs']), | |
| intensities = np.array(spec['formula_intensities']), | |
| metadata = {'precursor_mz': precursor_mz, 'formulas': np.array(spec['formulas'])}) | |
| if isinstance(self.spec_transform, dict): | |
| for key, transform in self.spec_transform.items(): | |
| item[key] = transform(spec) if transform is not None else spec | |
| else: | |
| item["spec"] = self.spec_transform(spec) | |
| else: | |
| item = self.instance.__getitem__(spec_i, transform_mol=False) | |
| item['cand'] = self.mol_transform(cand_smiles) | |
| item['cand_smiles'] = cand_smiles | |
| item['label'] = label | |
| if self.use_fp: | |
| item['fp'] = torch.Tensor(self.fpgen.GetFingerprint(Chem.MolFromSmiles(cand_smiles)).ToList()) | |
| return item | |
| class MassSpecDataset_Candidates: | |
| def __init__(self, | |
| use_formulas: bool, | |
| aug_cands_dir_pth: str, | |
| aug_cands_size: int, | |
| **kwargs): | |
| self.aug_cands_size = aug_cands_size | |
| self.instance = MassSpecDataset_PeakFormulas(**kwargs, return_mol_freq=False) if use_formulas else JESTR1_MassSpecDataset(**kwargs, return_mol_freq=False) | |
| with open(aug_cands_dir_pth, 'rb') as f: | |
| aug_cands = pickle.load(f) | |
| if self.use_fp: | |
| self.fpgen = AllChem.GetMorganGenerator(radius=5,fpSize=1024) | |
| self.aug_cands = {} | |
| targets = np.array(list(aug_cands.keys())) | |
| for smiles, cands in aug_cands.items(): | |
| # sort candidates by tanimoto similarity | |
| cands.sort(key=lambda x: x[1], reverse=True) | |
| cands = [c for c in cands if '.' not in c] | |
| # assert(len(cands) >0) | |
| if len(cands) <=1: # if no candidates, shuffle from target list | |
| np.random.shuffle(targets) | |
| cands = targets | |
| self.aug_cands[smiles] = itertools.cycle(cands) | |
| def __getattr__(self, name): | |
| return self.instance.__getattribute__(name) | |
| def __getitem__(self, i): | |
| item = self.instance.__getitem__(i,transform_mol=False) | |
| aug_cands = [next(self.aug_cands[item['mol']]) for _ in range(self.aug_cands_size)] | |
| item['aug_cands_fp'] = [self.fpgen.GetFingerprint(Chem.MolFromSmiles(c)).ToList() for c in aug_cands] | |
| item["aug_cands"] = [self.mol_transform(c) for c in aug_cands] | |
| item["mol"] = self.mol_transform(item["mol"]) | |
| return item |