changes

README.md

@@ -1,2 +1,9 @@
 # FILIP-MS
 FILIP contrastive learning for metabolite annotation
+
+## Magma preprocessing
+`python run_magma.py --data_pth '/r/hassounlab/spectra_data/msgym/MassSpecGym.tsv' --output_dir '/data/yzhouc01/FILIP-MS/data/magma/' --workers 50`
+
+### test candidates
+
+`python run_fragmentation_only.py --d '/r/hassounlab/spectra_data/msgym/MassSpecGym.tsv' -c '/r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_mass.json' -o '/data/yzhouc01/FILIP-MS/data/msgym_mass_cands.pkl' -w 50`

mvp/data/datasets.py

@@ -203,10 +203,10 @@ 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]
+        # 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)
@@ -369,6 +369,18 @@ class ExpandedRetrievalDataset:
             self.candidates[s] = [c for c in cand if '.' not in c]
         
         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()   
         
@@ -376,6 +388,7 @@ class ExpandedRetrievalDataset:
         
         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]
@@ -383,8 +396,8 @@ class ExpandedRetrievalDataset:
                 print(f"Skipping {spec_id}; empty candidate set")
                 continue
             if not any(labels):
-                print(f"Target smiles not in candidate set")
-
+                # 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)])
     

mvp/models/contrastive.py

@@ -245,6 +245,11 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
             self.result_dct[i]['candidates'].extend(cands)
             self.result_dct[i]['scores'].extend(scores.cpu().tolist())
             self.result_dct[i]['labels'].extend([x.cpu().item() for x in l])
+
+        # # external test case only
+        # for i, cands, scores in zip(outputs['identifiers'], outputs['cand_smiles'], outputs['scores']):
+        #     self.result_dct[i.cpu().item()]['candidates'].extend(cands)
+        #     self.result_dct[i.cpu().item()]['scores'].extend(scores.cpu().tolist())
             
     def _compute_rank(self, scores, labels):
         if not any(labels):
@@ -262,7 +267,6 @@ class ContrastiveModel(RetrievalMassSpecGymModel):
         self.df_test['rank'] = self.df_test.apply(lambda row: self._compute_rank(row['scores'], row['labels']), axis=1)
         if not self.df_test_path:
             self.df_test_path = os.path.join(self.hparams['experiment_dir'], 'result.pkl')
-        # self.df_test_path.parent.mkdir(parents=True, exist_ok=True)
         self.df_test.to_pickle(self.df_test_path)
 
     def get_checkpoint_monitors(self) -> T.List[dict]:

mvp/params_formSpec.yaml

@@ -1,6 +1,6 @@
 # Experiment setup 
 job_key: ''
-run_name: 'optimized_filip-model'
+run_name: 'simple_model'
 run_details: ""
 project_name: ''
 wandb_entity_name: 'mass-spec-ml'
@@ -10,7 +10,7 @@ debug: False
 checkpoint_pth:
 
 # Training setup
-max_epochs: 2000
+max_epochs: 1500
 accelerator: 'gpu'
 devices: [1]
 log_every_n_steps: 250
@@ -19,7 +19,7 @@ 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/MVP/data/sample/data.tsv #/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/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"
+subformula_dir_pth: /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default # /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default # /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: 
@@ -110,7 +110,7 @@ formula_dropout:  0.2
 formula_dims: [512, 256, 512] #[64, 128, 256]
 cross_attn_heads: 2
 use_cls: False
-peak_dropout: 0.414425691950033 # 0.2
+peak_dropout: 0.2
 formula_attn_heads: 4 # 2
 formula_transformer_layers: 2
 

mvp/run.sh

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

mvp/tune.py

@@ -116,19 +116,19 @@ def objective(trial: optuna.Trial, base_params, trial_times, base_dir, total_tri
 
     try:
         # Training-related params
-        params["batch_size"] = trial.suggest_categorical("batch_size", [64, 128])
+        params["batch_size"] = trial.suggest_categorical("batch_size", [32, 64, 128, 256])
         params["lr"] = trial.suggest_float("lr", 1e-6, 1e-3, log=True)
         params["weight_decay"] = trial.suggest_float("weight_decay", 1e-6, 1e-2, log=True)
-        params["contr_temp"] = trial.suggest_float("contrastive_temp", 0.02, 0.1)
+        params["contr_temp"] = trial.suggest_float("contrastive_temp", 0.01, 0.1)
 
         # Spectra encoder-related params
-        params['peak_dropout'] = trial.suggest_float("peak_dropout", 0.1, 0.5)
-        params['formula_attn_heads'] = trial.suggest_categorical("formula_attn_heads", [2, 4])
-        params['formula_transformer_layers'] = trial.suggest_categorical("formula_transformer_layers", [2, 4])
+        params['formula_dropout'] = trial.suggest_float("peak_dropout", 0.1, 0.5)
+        params['formula_attn_heads'] = trial.suggest_categorical("formula_attn_heads", [2, 4, 8])
+        params['formula_transformer_layers'] = trial.suggest_categorical("formula_transformer_layers", [1,2,3,4,5])
 
         choice = trial.suggest_categorical(
             "formula_dims",
-            ["64,128", "512,256", "256,512", "128", "256"]
+            ["64,128", "512,256", "256,512", "128", "256", "128,128", "512,512", "64,64,64,64"]
         )
         params["formula_dims"] = [int(x) for x in choice.split(",")]
 
@@ -136,12 +136,12 @@ def objective(trial: optuna.Trial, base_params, trial_times, base_dir, total_tri
         params['gnn_dropout'] = trial.suggest_float("gnn_dropout", 0.1, 0.5)
         choice = trial.suggest_categorical(
             "gnn_channels",
-            ["64,128", "128,256", "256,512", "64,128,128"]
+            ["64,128", "128,256", "256,512", "64,128,128", "128,128", "64,64,64"]
         )
         params["gnn_channels"] = [int(x) for x in choice.split(",")]
 
         # Ensure last layer matches final embedding dim
-        final_embedding_dim = trial.suggest_categorical("final_embedding_dim", [256, 512])
+        final_embedding_dim = trial.suggest_categorical("final_embedding_dim", [64,256,512,1024])
         params['formula_dims'].append(final_embedding_dim)
         params['gnn_channels'].append(final_embedding_dim)
 
@@ -229,8 +229,9 @@ def main(args):
     with open(args.param_pth) as f:
         params = yaml.load(f, Loader=yaml.FullLoader)
 
-    now = datetime.datetime.now().strftime("%Y%m%d")
-    base_dir = str(TEST_RESULTS_DIR / f"{now}_{params['run_name']}_optuna")
+    # now = datetime.datetime.now().strftime("%Y%m%d")
+    # base_dir = str(TEST_RESULTS_DIR / f"{now}_{params['run_name']}_optuna")
+    base_dir = "/data/yzhouc01/FILIP-MS/experiments/20250916_simple_model_optuna"
     os.makedirs(base_dir, exist_ok=True)
     params["experiment_dir"] = base_dir
 
@@ -239,8 +240,10 @@ def main(args):
     setup_logging(log_path)
 
     trial_times = []
+    study_name = "filip_contrastive"
+    storage = f"sqlite:///{base_dir}/optuna_study.db"
 
-    study = optuna.create_study(direction="minimize", pruner=optuna.pruners.MedianPruner())
+    study = optuna.create_study(study_name=study_name, storage=storage, direction="minimize", pruner=optuna.pruners.MedianPruner(), load_if_exists=True)
     study.optimize(lambda trial: objective(trial, params, trial_times, base_dir, args.n_trials), n_trials=args.n_trials)
 
     # Print best trial