magma loader

mvp/data/datasets.py

@@ -154,15 +154,13 @@ class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
         print("Data path: ", self.pth)
         self.metadata = pd.read_csv(self.pth, sep="\t")
 
-        # Used for training on consensus spectra
-        # with open(self.pth, 'rb') as f:
-        #     self.metadata = pickle.load(f)
-        # self.metadata['identifier'] = self.metadata['smiles'].tolist()
-
         # load subformulas
         all_spec_ids = self.metadata['identifier'].tolist()
         subformulaLoader = data_utils.Subformula_Loader(spectra_view=spectra_view, dir_path=subformula_dir_pth)
-        id_to_spec = subformulaLoader(all_spec_ids)
+        
+        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]

mvp/models/contrastive.py

@@ -270,6 +270,7 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
     def get_checkpoint_monitors(self) -> T.List[dict]:
         monitors = [
             {"monitor": f"{Stage.TRAIN.to_pref()}loss", "mode": "min", "early_stopping": False}, # monitor train loss
+            {"monitor": f"{Stage.VAL.to_pref()}loss", "mode": "min", "early_stopping": False}, # monitor val loss
         ]
         return monitors
     

mvp/params_formSpec.yaml

@@ -1,13 +1,13 @@
 # Experiment setup 
 job_key: ''
-run_name: 'filip_quick_test'
+run_name: 'sirius_labels'
 run_details: ""
 project_name: ''
 wandb_entity_name: 'mass-spec-ml'
 no_wandb: True
 seed: 0
 debug: False
-checkpoint_pth: #'../pretrained_models/msgym_formSpec.ckpt'
+checkpoint_pth:
 
 # Training setup
 max_epochs: 2000
@@ -19,10 +19,10 @@ 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: /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default #/data/yzhouc01/spectra_data/subformulae #"../data/MassSpecGym/data/subformulae_default"
+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"
 split_pth: 
-fp_dir_pth: '../data/MassSpecGym/data/morganfp_r5_1024.pickle'
-cons_spec_dir_pth: "../data/MassSpecGym/data/sample_consensus_formSpec.pkl"
+fp_dir_pth:
+cons_spec_dir_pth: 
 NL_spec_dir_pth: ""
 partial_checkpoint: ""
 

mvp/params_tmp.yaml

@@ -0,0 +1,121 @@
+# Experiment setup 
+job_key: ''
+run_name: 'filipContrastive'
+run_details: ""
+project_name: ''
+wandb_entity_name: 'mass-spec-ml'
+no_wandb: True
+seed: 0
+debug: False
+checkpoint_pth:
+
+# Training setup
+max_epochs: 2000
+accelerator: 'gpu'
+devices: [1]
+log_every_n_steps: 250
+val_check_interval: 1.0
+
+# Data paths
+candidates_pth: /data/yzhouc01/cancer/candidates.json # "../data/MassSpecGym/data/molecules/MassSpecGym_retrieval_candidates_formula.json"
+dataset_pth: /data/yzhouc01/cancer/data.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: /data/yzhouc01/cancer/subformulae # /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: 
+NL_spec_dir_pth: ""
+partial_checkpoint: ""
+
+# General hyperparameters
+batch_size: 64
+lr: 5.0e-05 
+weight_decay: 0
+contr_temp: 0.05
+early_stopping_patience: 300
+loss_strategy: 'static'
+num_workers: 50
+
+
+############################## Data transforms ##############################
+# - Spectra 
+spectra_view: SpecFormula #SpecMzIntTokens #SpecFormula 
+#  1. Binner
+max_mz: 1000
+bin_width: 1
+mask_peak_ratio: 0.00
+
+# 2. SpecFormula
+element_list: ['H', 'C',  'O', 'N', 'P', 'S', 'Cl', 'F', 'Br', 'I', 'B', 'As', 'Si', 'Se']
+add_intensities: True
+mask_precursor: False
+
+# - Molecule
+molecule_view: "MolGraph"
+atom_feature: 'full'
+bond_feature: 'full'
+
+
+############################## Views ##############################
+# contrastive
+use_contr: False
+contr_wt: 1
+contr_wt_update: {} 
+
+# consensus spectra
+use_cons_spec: False
+cons_spec_wt: 3
+cons_spec_wt_update: {} 
+cons_loss_type: 'l2' # cosine, l2
+
+# fp prediction/usage
+pred_fp: False
+use_fp: False
+fp_loss_type: 'cosine' #cosine, bce
+fp_wt: 3
+fp_wt_update: {} 
+fp_size: 1024
+fp_radius: 5
+fp_dropout: 0.4
+
+# candidates
+aug_cands: False
+aug_cands_wt: 0.1
+aug_cands_update: {} 
+aug_cands_size: 3
+
+# neutral loss
+use_NL: False
+
+
+############################## Task and model ##############################
+task: 'retrieval'
+spec_enc: Transformer_Formula # Transformer_MzInt #Transformer_Formula 
+mol_enc: "GNN"
+model:  filipContrastive # "MultiviewContrastive"
+contr_views: [['spec_enc', 'mol_enc']] #[['spec_enc', 'mol_enc'], ['spec_enc', 'NL_spec_enc'], ['mol_enc', 'NL_spec_enc']] #[['spec_enc', 'mol_enc'], ['mol_enc', 'cons_spec_enc'], ['cons_spec_enc', 'spec_enc'], ['fp_enc', 'mol_enc'], ['fp_enc', 'spec_enc'], ['fp_enc', 'cons_spec_enc']]
+log_only_loss_at_stages: []
+df_test_path: ""
+
+# - Spectra encoder 
+final_embedding_dim: 512
+fc_dropout: 0.4
+
+# - Spectra Token encoder
+hidden_dims: [64, 128]
+peak_dropout: 0.2
+
+# - Formula-based spec encoders
+formula_dropout: 0.2
+formula_dims: [64, 128, 256]
+cross_attn_heads: 2
+use_cls: False
+
+# -- GAT params
+attn_heads: [12,12,12]
+
+# - Molecule encoder (GNN)
+gnn_channels: [64,128,256] 
+gnn_type: "gcn"
+num_gnn_layers: 3
+gnn_hidden_dim: 512
+gnn_dropout: 0.3

mvp/run.sh

@@ -1,12 +1,3 @@
-# 1. preprocess data (subformula labels should be obtained through MIST)
-# python data_preprocess.py --spec_type formSpec --dataset_pth ../data/sample/data.tsv --candidates_pth  ../data/sample/candidates_mass.json --subformula_dir_pth ../data/sample/subformulae_default/ --output_dir ../data/sample/
-
-# 2. train model on msgym
-# python train.py --param_pth params_formSpec.yaml
-
-# 3. test model on msgym
-# python train.py --param_pth params_binnedSpec.yaml
-
-# python train.py
-python test.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

@@ -2,7 +2,7 @@ import argparse
 import datetime
 import sys
 sys.path.insert(0, "/data/yzhouc01/MassSpecGym")
-sys.path.insert(0, "/data/yzhouc01/MVP")
+sys.path.insert(0, "/data/yzhouc01/FILIP-MS")
 
 from rdkit import RDLogger
 import pytorch_lightning as pl

mvp/train.py

@@ -33,7 +33,7 @@ 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['candidates_pth'] =None
         params['split_pth']=None
 
@@ -80,7 +80,7 @@ def main(params):
             filename=f'{{epoch}}-{{{monitor_name}:.2f}}',
             # filename='{epoch}-{val_loss:.2f}-{train_loss:.2f}',
             auto_insert_metric_name=True,
-            save_last=(i == 0)
+            # save_last=(i == 0)
         )
         callbacks.append(checkpoint)
         if monitor.get('early_stopping', False):

mvp/utils/data.py

@@ -9,11 +9,18 @@ import mvp.data.datasets as jestr_datasets
 import typing as T
 from mvp.definitions import MSGYM_FORMULA_VECTOR_NORM, MSGYM_STANDARD_MH
 import matchms
+import tqdm
 
 class Subformula_Loader:
-    def __init__(self, spectra_view, dir_path) -> None:
+    """
+    :param dir_path: path to folder containing either MIST or SIRIUS formulas, automatically parses the file type as needed
+    :param use_prec_mz: add precursor m/z when fragment precursor peak is not present or remove precursor peak when their is no fragment precursor peak
+    """
+    def __init__(self, spectra_view, dir_path, use_prec_mz=True, formula_source='default') -> None:
 
         self.dir_path = dir_path
+        self.use_prec_mz = use_prec_mz
+        self.formula_source = formula_source
         if spectra_view == 'SpecFormula':
             self.load = self.load_subformula_data
         elif spectra_view == "SpecFormulaMz":
@@ -21,37 +28,90 @@ class Subformula_Loader:
         else:
             raise Exception("Spectra view is not supported.")
 
-    def __call__(self, ids):
+    def __call__(self, ids, form_list, prec_mz_list):
         id_to_form_spec = {}
-        for id in ids:
-            data = self.load(id)
-            if data:
+        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:
                 id_to_form_spec[id] = data
-        
+
         return id_to_form_spec
-    
-    def load_subformula_data(self, spec_id: str):
-        '''MIST subformula format:https://github.com/samgoldman97/mist/blob/main_v2/src/mist/utils/spectra_utils.py 
+
+    def load_mist_data(self, data, curr_form, curr_prec_mz):
+        '''MIST subformula format:https://github.com/samgoldman97/mist/blob/main_v2/src/mist/utils/spectra_utils.py
         '''
         try:
-            file = os.path.join(self.dir_path, spec_id+".json")
-            with open(file) as f:
-                data = json.load(f)
+            # file = os.path.join(self.dir_path, spec_id+".json")
+            # with open(file) as f:
+            #     data = json.load(f)
             mzs = np.array(data['output_tbl']['mz'])
             formulas = np.array(data['output_tbl']['formula'])
             intensities = np.array(data['output_tbl']['ms2_inten'])
 
+            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]])
+            elif curr_form in formulas and self.use_prec_mz:
+                idx = np.where(formulas == curr_form)[0][0]
+                intensities[idx] = 1.1
+
+            # sort by mzs
+            ind = mzs.argsort()
+            mzs = mzs[ind]
+            formulas = formulas[ind]
+            intensities = intensities[ind]
+            return {'formulas': formulas, 'formula_mzs': mzs, 'formula_intensities': intensities}
+        except:
+            return None
+
+    def load_magma_data(self, data, curr_form, curr_prec_mz):
+
+        return None
+
+    def load_sirius_data(self, data):
+        try:
+
+            mzs = np.array([entry['mz'] for entry in data['fragments']])
+            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
+
+            if not self.use_prec_mz:
+                not_append_prec_mz = np.array([len(entry['peaks']) != 0 for entry in data['fragments']])
+
+                mzs = mzs[not_append_prec_mz]
+                formulas = formulas[not_append_prec_mz]
+                intensities = intensities[not_append_prec_mz]
+
             # sort by mzs
             ind = mzs.argsort()
             mzs = mzs[ind]
             formulas = formulas[ind]
             intensities = intensities[ind]
             return {'formulas': formulas, 'formula_mzs': mzs, 'formula_intensities': intensities}
+        except:
+            
+            return None
+
+    def load_subformula_data(self, spec_id: str, curr_form: str, curr_prec_mz: float):
+        try:
+            file = os.path.join(self.dir_path, spec_id+".json")
+            with open(file) as f:
+                data = json.load(f)
+            if self.formula_source == 'sirius':
+                return self.load_sirius_data(data)
+            elif self.formula_source == 'magma':
+                return self.load_magma_data(data, curr_form, curr_prec_mz)
+            else:
+                return self.load_mist_data(data, curr_form, curr_prec_mz)
+
         except:
             return None
 
     def load_subformula_dict(self, spec_id: str):
-        '''MIST subformula format:https://github.com/samgoldman97/mist/blob/main_v2/src/mist/utils/spectra_utils.py 
+        '''MIST subformula format:https://github.com/samgoldman97/mist/blob/main_v2/src/mist/utils/spectra_utils.py
         '''
         try:
             file = os.path.join(self.dir_path, spec_id+".json")
@@ -64,7 +124,7 @@ class Subformula_Loader:
             mz_to_formulas = {mz:f for mz, f in zip(mzs, formulas)}
             for mz, f in zip(mzs, formulas):
                 mz_to_formulas[mz] = f
-            
+
             ind = mzs.argsort()
             mzs = mzs[ind]
             formulas = formulas[ind]