update

mvp/data/datasets.py

@@ -19,6 +19,9 @@ 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,
@@ -90,8 +93,6 @@ class JESTR1_MassSpecDataset(MassSpecDataset):
                     item[key] = transform(spec) if transform is not None else spec
             else:
                 item["spec"] = self.spec_transform(spec)
-        else:
-            item["spec"] = spec
 
         if self.return_mol_freq:
             item["mol_freq"] = metadata["mol_freq"]
@@ -132,7 +133,9 @@ class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
         cons_spec_dir_pth: str = None,
         return_mol_freq: bool = False,
         return_identifier: bool = True,
-        dtype: T.Type = torch.float32
+        dtype: T.Type = torch.float32,
+        formula_source = 'default',
+        stage: Stage = Stage.TRAIN
     ):
         """
         Args:
@@ -146,6 +149,8 @@ class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
         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)
@@ -155,19 +160,7 @@ class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
         self.metadata = pd.read_csv(self.pth, sep="\t")
 
         # load subformulas
-        all_spec_ids = self.metadata['identifier'].tolist()
-        subformulaLoader = data_utils.Subformula_Loader(spectra_view=spectra_view, dir_path=subformula_dir_pth)
-        
-        form_list = self.metadata['formula'].tolist()
-        prec_mz_list = self.metadata['precursor_mz'].tolist()
-        id_to_spec = 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())))
-        
+        id_to_spec = self._load_id_to_spec(stage)
         
         # load fingerprints
         self._load_fp(fp_dir_pth)
@@ -179,6 +172,7 @@ class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
         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')
 
@@ -208,6 +202,27 @@ class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
 
         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,
@@ -255,7 +270,11 @@ class ContrastiveDataset(Dataset):
         # standard collate
         for k in batch[0].keys():
             if k not in non_standard_collate:
-                collated_batch[k] = default_collate([item[k] for item in batch])
+                try:
+                    collated_batch[k] = default_collate([item[k] for item in batch])
+                except:
+                    print(f"Error in collating key {k}")
+                    raise
                 
         # batch graphs
         batch_mol = []
@@ -327,10 +346,13 @@ class ExpandedRetrievalDataset:
                  candidates_pth: T.Optional[T.Union[Path, str]] = None,
                  fp_size: int = None,
                  fp_radius: int = None,
+                 use_magma = False,
                 **kwargs):
+
         
-        self.instance = MassSpecDataset_PeakFormulas(**kwargs, return_mol_freq=False) if use_formulas else JESTR1_MassSpecDataset(**kwargs, return_mol_freq=False)
-        # super().__init__(**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)
@@ -348,9 +370,10 @@ class ExpandedRetrievalDataset:
         
         self.spec_cand = [] #(spec index, cand_smiles, true_label)
         test_smiles = self.metadata[self.metadata['fold'] == "test"]['smiles'].tolist()
-        test_ms_id = self.metadata[self.metadata['fold'] == "test"]['identifier'].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))
         
-        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)
@@ -363,7 +386,7 @@ class ExpandedRetrievalDataset:
                 print(f"Target smiles not in candidate set")
 
 
-            self.spec_cand.extend([(spec_id_to_index[spec_id], candidates[j], k) for j, k in enumerate(labels)])
+            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)
@@ -376,7 +399,33 @@ class ExpandedRetrievalDataset:
         cand_smiles = self.spec_cand[i][1]
         label = self.spec_cand[i][2]
 
-        item = self.instance.__getitem__(spec_i, transform_mol=False)
+        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

mvp/data/transforms.py

@@ -160,7 +160,7 @@ class SpecFormulaMzFeaturizer(SpecTransform):
                         # print(f"Couldn't vectorize {f}, element {e} not supported")
                         continue
         return formula_vector
-    
+
 class SpecFormulaFeaturizer(SpecTransform):
     ''' Uses processed mz and intensities, excludes mz values, keep peaks with formulas only'''
     def __init__(
@@ -208,7 +208,7 @@ class SpecFormulaFeaturizer(SpecTransform):
                     try:
                         formula_vector[i][self.elem_to_pos[e]]+=ct
                     except:
-                            print(f"Couldn't vectorize {f}, element {e} not supported")
+                            # print(f"Couldn't vectorize {f}, element {e} not supported")
                             continue
             except:
                 print(f"Couldn't vectorize {f}, formula not supported")

mvp/definitions.py

@@ -40,4 +40,6 @@ MSGYM_STANDARD_MH = {
 }
 MSGYM_STANDARD_all = { # got these from Yinkai
 "mz_mean": 80.88304948022557,
-"mz_std" : 197.4588028571758}
+"mz_std" : 197.4588028571758}
+
+PRECURSOR_INTENSITY = 1.1

mvp/params_formSpec.yaml

@@ -1,6 +1,6 @@
 # Experiment setup 
 job_key: ''
-run_name: 'sirius_labels'
+run_name: 'magma_all_labels'
 run_details: ""
 project_name: ''
 wandb_entity_name: 'mass-spec-ml'
@@ -12,14 +12,14 @@ checkpoint_pth:
 # Training setup
 max_epochs: 2000
 accelerator: 'gpu'
-devices: [1]
+devices: [0]
 log_every_n_steps: 250
 val_check_interval: 1.0
 
 # Data paths
 candidates_pth: /r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_mass.json # "../data/MassSpecGym/data/molecules/MassSpecGym_retrieval_candidates_formula.json"
 dataset_pth: /r/hassounlab/spectra_data/msgym/MassSpecGym.tsv #/data/yzhouc01/spectra_data/combined_msgym_nist23_multiplex.tsv # /r/hassounlab/spectra_data/msgym/MassSpecGym.tsv # "../data/MassSpecGym/data/sample_data.tsv" 
-subformula_dir_pth: /r/hassounlab/msgym_sirius # /r/hassounlab/msgym_sirius # /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default #/data/yzhouc01/spectra_data/subformulae #"../data/MassSpecGym/data/subformulae_default"
+subformula_dir_pth: /data/yzhouc01/FILIP-MS/data/magma # /r/hassounlab/msgym_sirius # /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default #/data/yzhouc01/spectra_data/subformulae #"../data/MassSpecGym/data/subformulae_default"
 split_pth: 
 fp_dir_pth:
 cons_spec_dir_pth: 
@@ -39,6 +39,7 @@ num_workers: 50
 ############################## Data transforms ##############################
 # - Spectra 
 spectra_view: SpecFormula #SpecMzIntTokens #SpecFormula 
+formula_source: 'magma_all' # magma_1, magma_all, sirius, default
 #  1. Binner
 max_mz: 1000
 bin_width: 1

mvp/run.sh

@@ -1,3 +1,3 @@
-python train.py
+# python train.py
 python test.py 
 python test.py --candidates_pth /r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_formula.json

mvp/test.py

@@ -35,12 +35,14 @@ def main(params):
         
     # Init paths to data files
     if params['debug']:
-        params['dataset_pth'] = "../data/sample/data.tsv"
+
+        params['dataset_pth'] = "/data/yzhouc01/MVP/data/sample/data.tsv"
         params['split_pth']=None
         params['df_test_path'] = os.path.join(params['experiment_dir'], 'debug_result.pkl')
 
     # Load dataset
     spec_featurizer = get_spec_featurizer(params['spectra_view'], params)
+
     mol_featurizer = get_mol_featurizer(params['molecule_view'], params)
     dataset = get_test_ms_dataset(params['spectra_view'], params['molecule_view'], spec_featurizer, mol_featurizer, params)
 

mvp/utils/data.py

@@ -7,7 +7,7 @@ from massspecgym.data.transforms import SpecTransform, MolTransform
 from mvp.data.transforms import MolToGraph
 import mvp.data.datasets as jestr_datasets
 import typing as T
-from mvp.definitions import MSGYM_FORMULA_VECTOR_NORM, MSGYM_STANDARD_MH
+from mvp.definitions import MSGYM_FORMULA_VECTOR_NORM, MSGYM_STANDARD_MH, PRECURSOR_INTENSITY
 import matchms
 import tqdm
 
@@ -30,6 +30,7 @@ class Subformula_Loader:
 
     def __call__(self, ids, form_list, prec_mz_list):
         id_to_form_spec = {}
+        print("Processing formula spectra")
         for id, curr_form, curr_prec_mz in tqdm.tqdm(zip(ids, form_list, prec_mz_list), total=len(ids)):
             data = self.load(id, curr_form, curr_prec_mz)
             if data is not None:
@@ -51,10 +52,10 @@ class Subformula_Loader:
             if curr_form not in formulas and self.use_prec_mz:
                 mzs = np.concatenate([mzs, [curr_prec_mz]])
                 formulas = np.concatenate([formulas, [curr_form]])
-                intensities = np.concatenate([intensities, [1.1]])
+                intensities = np.concatenate([intensities, [PRECURSOR_INTENSITY]])
             elif curr_form in formulas and self.use_prec_mz:
                 idx = np.where(formulas == curr_form)[0][0]
-                intensities[idx] = 1.1
+                intensities[idx] = PRECURSOR_INTENSITY
 
             # sort by mzs
             ind = mzs.argsort()
@@ -66,8 +67,75 @@ class Subformula_Loader:
             return None
 
     def load_magma_data(self, data, curr_form, curr_prec_mz):
+        
+        np.random.seed(42)
+
+        formula_to_intensity = {}
+        formula_to_mz = {}
+
+        # data is None
+        if data is None:
+            if self.use_prec_mz:
+                return {'formulas': [curr_form], 'formula_mzs': [curr_prec_mz], 'formula_intensities': [PRECURSOR_INTENSITY]}
+            else:
+                return {'formulas': [], 'formula_mzs': [], 'formula_intensities': []}
+
+        # randomly choose 1 formula for each peak, keep largest intensity for each formula
+        if self.formula_source.endswith('1'):
+            for f, m, i in zip(data['subformulas'], data['mz'], data['intensities']):
+
+                if not f:
+                    continue 
+                selected_f = np.random.choice(f)
+                if selected_f in formula_to_intensity:
+                    if i > formula_to_intensity[selected_f]:
+                        formula_to_intensity[selected_f] = i
+                        formula_to_mz[selected_f] = m
+                else:
+                    formula_to_intensity[selected_f] = i
+                    formula_to_mz[selected_f] = m
+
+        # take all formulas, divide intensity by number of formulas, keep largest intensity for each formula
+        elif self.formula_source.endswith('all'):
+            for f, m, i in zip(data['subformulas'], data['mz'], data['intensities']):
+
+                if not f:
+                    continue
+                for fi in f:
+                    if fi in formula_to_intensity:
+                        if i/len(f) > formula_to_intensity[fi]:
+                            formula_to_intensity[fi] = i/len(f)
+                            formula_to_mz[fi] = m
+                    else:
+                        formula_to_intensity[fi] = i/len(f)
+                        formula_to_mz[fi] = m
+        else:
+            raise Exception(f"Formula source not supported: {self.formula_source}")
+
+        mzs = list(formula_to_mz.values())
+        formulas = list(formula_to_mz.keys())
+        intensities = list(formula_to_intensity.values())
+
+        # add precursor mz
+        if self.use_prec_mz:
+            if curr_form in formulas:
+                intensities[formulas.index(curr_form)] = PRECURSOR_INTENSITY
+            else:
+                formulas.append(curr_form)
+                intensities.append(PRECURSOR_INTENSITY)
+                mzs.append(curr_prec_mz)
+
+        # sort by mzs
+        mzs = np.array(mzs)
+        formulas = np.array(formulas)
+        intensities = np.array(intensities)
+
+        ind = mzs.argsort()
+        mzs = mzs[ind]
+        formulas = formulas[ind]
+        intensities = intensities[ind]
 
-        return None
+        return {'formulas': formulas, 'formula_mzs': mzs, 'formula_intensities': intensities}
 
     def load_sirius_data(self, data):
         try:
@@ -76,9 +144,9 @@ class Subformula_Loader:
             formulas = np.array([entry['molecularFormula'] for entry in data['fragments']])
             intensities = np.array([entry['relativeIntensity'] for entry in data['fragments'] ])
 
-            intensities[formulas == data['molecularFormula']] = 1.1
+            intensities[formulas == data['molecularFormula']] = PRECURSOR_INTENSITY
 
-            if not self.use_prec_mz:
+            if not self.use_prec_mz: # removing precursor formula
                 not_append_prec_mz = np.array([len(entry['peaks']) != 0 for entry in data['fragments']])
 
                 mzs = mzs[not_append_prec_mz]
@@ -102,7 +170,7 @@ class Subformula_Loader:
                 data = json.load(f)
             if self.formula_source == 'sirius':
                 return self.load_sirius_data(data)
-            elif self.formula_source == 'magma':
+            elif self.formula_source.startswith('magma'):
                 return self.load_magma_data(data, curr_form, curr_prec_mz)
             else:
                 return self.load_mist_data(data, curr_form, curr_prec_mz)
@@ -200,7 +268,7 @@ def get_test_ms_dataset(spectra_view: T.Union[str, T.List[str]],
 
     dataset_params = {'spectra_view': spectra_view, 'pth': params['dataset_pth'], 'spec_transform': spectra_featurizer, 'mol_transform': mol_featurizer, "candidates_pth": params['candidates_pth']}
     if "SpecFormula" in views or "SpecFormulaMz" in views:
-        dataset_params.update({'subformula_dir_pth': params['subformula_dir_pth']})
+        dataset_params.update({'subformula_dir_pth': params['subformula_dir_pth'], 'use_magma': params['formula_source'].startswith('magma'), 'formula_source':params['formula_source']})
         use_formulas = True 
     
     if params['use_cons_spec']:
@@ -223,7 +291,7 @@ def get_ms_dataset(spectra_view: str,
     dataset_params = {'pth': params['dataset_pth'], 'spec_transform': spectra_featurizer, 'mol_transform': mol_featurizer, 'spectra_view': spectra_view}
     use_formulas = False
     if "SpecFormula" in spectra_view:
-        dataset_params.update({'subformula_dir_pth': params['subformula_dir_pth']})
+        dataset_params.update({'subformula_dir_pth': params['subformula_dir_pth'], 'formula_source': params['formula_source']})
         use_formulas = True
 
     if params['pred_fp'] or params['use_fp']: