clean up

README.md

@@ -8,55 +8,158 @@ pinned: false
 python_version: 3.11.7
 ---
 
-# 🔥 FLARE
-Fine-grained Learning for Aligment of spectra-molecule REpresentation
+# FLARE
+
+**F**ine-grained **L**earning for **A**lignment of spectra–molecule **RE**presentations
+
+### Authors
 
-### Authors  
 **Yan Zhou Chen, Soha Hassoun**  
-Department of Computer Science, Tufts University  
+Department of Computer Science, Tufts University
+
+---
+
+## Overview
+
+FLARE learns a joint embedding space for **MS/MS spectra** (represented as **per-peak chemical formulas** from a subformula assigner) and **molecular graphs**. The default publication model uses **FILIP-style contrastive learning** (`filipContrastive`): fine-grained similarity between spectrum tokens and graph nodes, with a temperature-scaled loss.
+
+Use cases:
+
+- **Retrieval**: rank a list of candidate SMILES for each query spectrum (MassSpecGym-style evaluation).
+- **Interpretation**: the Streamlit app visualizes **peak-to-node** correspondence for a single spectrum–molecule pair.
 
 ---
 
-FLARE is a framework for **ranking molecular candidates given a mass spectrum**. Beyond candidate ranking, FLARE provides **visualization of peak-to-node attribution**, enabling deeper insights into how spectral peaks correspond to molecular graph nodes.
+## Model (default stack)
+
+| Component | Setting (see `params.yaml`) |
+|-----------|----------------------------|
+| Spectrum input | `SpecFormula` — formula peaks from JSON in `subformula_dir_pth` |
+| Formula source | `default` — MIST-compatible JSON (`load_mist_data`); optional `sirius` |
+| Spectrum encoder | `Transformer_Formula` |
+| Molecule encoder | `GNN` (DGL + dgllife GCN), node embeddings for FILIP |
+| Training objective | `filipContrastive` — masked FILIP loss, temperature `contr_temp` |
+| Output | Embeddings for cosine / FILIP similarity at test time |
+
+Hyperparameters are split into: **run/logging**, **training loop**, **data paths**, **featurizers**, **encoder widths/depths**, and **evaluation** (`at_ks`, `myopic_mces_kwargs`). Only keys present in `params.yaml` are required; paths can be **relative to the repository root** (recommended) or absolute.
 
 ---
 
-## 🌐 Visualize Peak-to-Node Correspondence
-Explore our interactive [app](https://huggingface.co/spaces/HassounLab/FLARE) to visualize peak-to-node attributes in real time.
+## Repository layout
+
+| Path | Role |
+|------|------|
+| `params.yaml` | Canonical training/testing/app hyperparameters |
+| `hparams.yaml` | Symlink to `params.yaml` (Hugging Face Spaces convention) |
+| `flare/` | Training (`train.py`, `test.py`, `tune.py`), models, data pipeline |
+| `massspecgym/` | Vendored MassSpecGym Lightning base classes and utilities |
+| `app.py`, `app_utils/` | Streamlit peak–node visualization |
+| `pretrained_models/` | Place public checkpoints here (e.g. `flare.ckpt`) |
+| `experiments/` | Default output root for new runs (see `flare/definitions.py`) |
+| `archive/` | Older scripts and features **not** part of the slim release (MAGMA, class experiments, legacy YAML, etc.); nothing was deleted |
 
 ---
-## 🛠 Set up
-### Clone repository
-```
+
+## Environment variables (no hardcoded machine paths)
+
+| Variable | Purpose |
+|----------|---------|
+| `FLARE_PARAMS` | Path to YAML params (default: `<repo>/params.yaml`) |
+| `FLARE_CHECKPOINT` | Checkpoint for the app or manual runs |
+| `FLARE_DEBUG_DATASET` | When `debug: true`, TSV path for a tiny local dataset |
+| `FLARE_REPO_ROOT` | Optional; overrides repo root for resolving relative paths in `default_param_path()` |
+| `MASSSPECGYM_ROOT` | Optional extra `sys.path` root if you use an external `massspecgym` checkout |
+| `FLARE_UPLOAD_CKPT`, `HF_REPO_ID`, `HF_REPO_TYPE`, `HF_TOKEN` | See `app_utils/upload_model.py` for HF uploads |
+
+---
+
+## Setup
+
+```bash
 git clone https://huggingface.co/spaces/HassounLab/FLARE
-cd flare
-```
-### Set up environment and install dependencies
-```
+cd FLARE
+
 conda create -n flare python=3.11
 conda activate flare
 pip install -r requirements.txt
 ```
+
+Place **MassSpecGym** (or your) spectrum TSV, **candidate JSON**, and **subformula JSON directory** where you want them, then set paths in `params.yaml` (relative paths like `data/MassSpecGym.tsv` resolve from the repo root).
+
+---
+
+## Data preparation
+
+Per-spectrum subformula JSON files (one file per spectrum id, MIST-style) are required for `SpecFormula`. Generate them with the bundled assigner (adapted from MIST):
+
+```bash
+cd flare/subformula_assign
+export SPEC_FILES=/path/to/spectra.tsv
+export OUTPUT_DIR=/path/to/subformulae_out
+export LABELS_FILE=/path/to/spectra.tsv   # often same as SPEC_FILES
+export MAX_FORMULAE=60
+bash run.sh
+```
+
+Defaults in `run.sh` point at `data/sample/` under the repo if you add a small sample there.
+
+---
+
+## Training
+
+From the repository root (so `flare` and `massspecgym` import correctly):
+
+```bash
+cd flare
+python train.py                          # uses FLARE_PARAMS or ../params.yaml
+python train.py --param_pth /path/to/custom.yaml
+```
+
+`train.py` creates `experiments/<YYYYMMDD>_<run_name>/`, writes TensorBoard logs there, and saves checkpoints. `df_test_path` defaults to `<experiment_dir>/result.pkl` if unset.
+
+---
+
+## Testing (retrieval)
+
+```bash
+cd flare
+python test.py \
+  --checkpoint_pth /path/to/epoch=....ckpt \
+  --exp_dir /path/to/experiment_dir   # optional; else latest matching run_name
+```
+
+Useful flags: `--candidates_pth`, `--df_test_pth`, `--external_test` (no positive label in the list). Override params file with `--param_pth` or `FLARE_PARAMS`.
+
 ---
-## 🚀 Usage
-Modify params.yaml as necessary
 
+## Hyperparameter search
+
+```bash
+cd flare
+python tune.py --n_trials 20
 ```
-# preprocess data
-python subformula_assign/assign_subformulae.py --spec-files ../data/sample/data.tsv --output-dir ../data/sample/subformulae --labels-file ../data/sample/data.tsv --max-formulae 60
 
-# train 
-python train.py
+Uses Optuna; study database and logs live under `experiments/<date>_<run_name>_optuna/`. Best YAML is written to `best_params.yaml` in that folder.
+
+---
+
+## Streamlit app (peak-to-node visualization)
 
-# test
-python test.py
+```bash
+streamlit run app.py
 ```
 
+The app loads architecture settings from `FLARE_PARAMS` (default `params.yaml`) and weights from `FLARE_CHECKPOINT` (default `pretrained_models/flare.ckpt`). Ensure the checkpoint matches the architecture in the YAML.
+
 ---
-## 🙏 Acknowledgments
-- **Training Data**: [MassSpecGym](https://github.com/pluskal-lab/MassSpecGym)
-- **Subformula Assigner Code**: [MIST](https://github.com/samgoldman97/mist/tree/main_v2)
+
+## Acknowledgments
+
+- **Data**: [MassSpecGym](https://github.com/pluskal-lab/MassSpecGym)
+- **Subformula tooling**: [MIST](https://github.com/samgoldman97/mist/tree/main_v2)
 
 ---
-## 📧 Contact
-For questions, reach out to soha.hassoun@tufts.edu
+
+## Contact
+
+For questions: soha.hassoun@tufts.edu

flare/data/datasets.py

@@ -1,124 +1,30 @@
-import pandas as pd
 import json
 import typing as T
+from collections import defaultdict
+from pathlib import Path
+
+import dgl
 import numpy as np
+import pandas as pd
 import torch
+import flare.utils.data as data_utils
 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
+import matchms
 from massspecgym.data.datasets import MassSpecDataset
-import flare.utils.data as data_utils
-from torch.nn.utils.rnn import pad_sequence
+from massspecgym.data.transforms import MolTransform, MolToInChIKey, SpecTransform
 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
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data.dataloader import default_collate
+from torch.utils.data.dataset import Dataset
+
 
 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)
+    """Same as MassSpecDataset; keeps `spectra_view` for API compatibility."""
 
-        self.use_fp = False
-        self.use_cons_spec = False
-        self.use_NL_spec = False
+    def __init__(self, spectra_view: str, **kwargs):
+        super().__init__(**kwargs)
         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__(
@@ -128,26 +34,16 @@ class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
         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
+        formula_source: str = "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
@@ -155,31 +51,23 @@ class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
         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)]
 
-        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")
 
-        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)
-                
+        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)
@@ -187,108 +75,104 @@ class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
 
         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
+        item = super().__getitem__(i, transform_spec, transform_mol=False)
+        mol = item["mol"]
 
-        # 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)
+        else:
+            item["mol"] = 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()
+        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())))
+        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,
-    ):
+    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_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:
+
+    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
+        self.smiles_to_spec_couter[mol] = counter + 1
         return item
 
     @staticmethod
-    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']
+    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"]
         require_pad = False
-        if 'Formula' in spectra_view or 'Tokens' in spectra_view:
+        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
+            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')
+            non_standard_collate.remove("aug_cands")
 
         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:
+                except Exception:
                     print(f"Error in collating key {k}")
                     raise
-                
-        # batch graphs
+
         if batch_mol:
-            batch_mol = []
-            batch_mol_nodes= []
+            batch_mol_list = []
+            batch_mol_nodes = []
 
             for item in batch:
-                batch_mol.append(item[mol_key])
+                batch_mol_list.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
+            collated_batch[mol_key] = dgl.batch(batch_mol_list)
+            collated_batch["mol_n_nodes"] = batch_mol_nodes
+
         if require_pad:
             peaks = []
             n_peaks = []
@@ -296,54 +180,40 @@ class ContrastiveDataset(Dataset):
                 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
+            collated_batch["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)
+    """Test-time retrieval over a fixed candidate pool per spectrum/formula."""
+
+    def __init__(
+        self,
+        use_formulas: bool = True,
+        mol_label_transform: MolTransform = MolToInChIKey(),
+        candidates_pth: T.Optional[T.Union[Path, str]] = None,
+        formula_to_smiles_pth: T.Optional[T.Union[Path, str]] = None,
+        external_test: bool = False,
+        **kwargs,
+    ):
+        self.external_test = external_test
+
+        self.instance = (
+            MassSpecDataset_PeakFormulas(**kwargs, return_mol_freq=False, stage=Stage.TEST)
+            if use_formulas
+            else JESTR1_MassSpecDataset(**kwargs, return_mol_freq=False)
+        )
 
         self.candidates_pth = candidates_pth
+        self.formula_to_smiles_pth = formula_to_smiles_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:
+
+        candidate_source_pth = self.formula_to_smiles_pth if self.formula_to_smiles_pth else self.candidates_pth
+        if not candidate_source_pth:
+            raise ValueError("One of candidates_pth or formula_to_smiles_pth must be provided.")
+
+        with open(candidate_source_pth, "r") as file:
             candidates = json.load(file)
 
         self.candidates = {}
@@ -351,130 +221,66 @@ class ExpandedRetrievalDataset:
             clean_cands = []
             for c in cand:
                 try:
-                    if '.' not in c:
+                    if "." not in c:
                         clean_cands.append(c)
-                except:
+                except Exception:
                     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]
+            self.candidates[s] = clean_cands
 
-        if self.use_magma:
-            item = self.instance.__getitem__(spec_i, transform_mol=False, transform_spec=False)
+        self.spec_cand = []
 
-            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)
+        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"]
 
-            spec = self.subformulaLoader.load_magma_data(spec_data, formula, precursor_mz)
+        if self.formula_to_smiles_pth:
+            if "formula" not in self.metadata.columns:
+                raise ValueError("formula_to_smiles_pth was provided, but dataset has no 'formula' column.")
+            self.metadata["candidate_key"] = self.metadata["formula"].astype(str)
+        else:
+            self.metadata["candidate_key"] = self.metadata["smiles"].astype(str)
 
-            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)
+        self.metadata = self.metadata[self.metadata["candidate_key"].isin(self.candidates.keys())]
 
+        if "fold" in self.metadata.columns:
+            test_metadata = self.metadata[self.metadata["fold"] == "test"]
         else:
-            item = self.instance.__getitem__(spec_i, transform_mol=False)
+            test_metadata = self.metadata
 
-        item['cand'] = self.mol_transform(cand_smiles)
-        item['cand_smiles'] = cand_smiles
-        item['label'] = label
+        self.spec_id_to_index = dict(zip(self.metadata["identifier"], self.metadata.index))
 
-        if self.use_fp:
-            item['fp'] = torch.Tensor(self.fpgen.GetFingerprint(Chem.MolFromSmiles(cand_smiles)).ToList())
+        for _, row in test_metadata.iterrows():
+            spec_id = row["identifier"]
+            candidate_key = row["candidate_key"]
+            cands = self.candidates[candidate_key]
 
-        return item
+            if self.external_test:
+                labels = [False for _ in cands]
+            else:
+                target_smiles = row["smiles"]
+                labels = [c == target_smiles for c in cands]
+            if len(cands) == 0:
+                print(f"Skipping {spec_id}; empty candidate set")
+                continue
 
-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)
+            self.spec_cand.extend([(self.spec_id_to_index[spec_id], cands[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):
-        item = self.instance.__getitem__(i,transform_mol=False)
+        spec_i = self.spec_cand[i][0]
+        cand_smiles = self.spec_cand[i][1]
+        label = self.spec_cand[i][2]
 
-        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"])
+        item = self.instance.__getitem__(spec_i, transform_mol=False)
+        item["cand"] = self.mol_transform(cand_smiles)
+        item["cand_smiles"] = cand_smiles
+        item["label"] = label
 
-        return item
+        return item

flare/run.sh

@@ -1,3 +1,17 @@
-# python train.py --param_pth params_filipGlobal.yaml
-# python test.py --param_pth params_filipGlobal.yaml
-python test.py --param_pth params_filipGlobal.yaml --candidates_pth /r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_formula.json
+#!/usr/bin/env bash
+# Example: run from the `flare/` directory.
+#   conda activate flare
+#   python train.py
+#   python test.py --checkpoint_choice val
+#
+# Optional overrides:
+#   export FLARE_PARAMS=/path/to/params.yaml
+#   export CANDIDATES_JSON="$PWD/../data/MassSpecGym_retrieval_candidates_formula.json"
+#   python test.py --candidates_pth "${CANDIDATES_JSON}"
+
+set -euo pipefail
+REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+export FLARE_REPO_ROOT="${FLARE_REPO_ROOT:-$REPO_ROOT}"
+
+# python train.py
+# python test.py --checkpoint_choice val

flare/test.py

@@ -1,128 +1,163 @@
 import argparse
 import datetime
+import os
 import sys
-sys.path.insert(0, "/data/yzhouc01/MassSpecGym")
-sys.path.insert(0, "/data/yzhouc01/FILIP-MS")
+from pathlib import Path
+
+
+def _add_local_dependency_paths() -> None:
+    flare_repo_root = Path(__file__).resolve().parents[1]
+    candidate_roots = [flare_repo_root, flare_repo_root / "massspecgym"]
+    for env_var in ("MASSSPECGYM_ROOT",):
+        env_value = os.environ.get(env_var)
+        if env_value:
+            candidate_roots.append(Path(env_value).expanduser())
+    for root in candidate_roots:
+        if root.exists():
+            root_str = str(root)
+            if root_str not in sys.path:
+                sys.path.insert(0, root_str)
+
+
+_add_local_dependency_paths()
+
+from functools import partial
 
-from rdkit import RDLogger
 import pytorch_lightning as pl
 from pytorch_lightning import Trainer
+from rdkit import RDLogger
 from massspecgym.models.base import Stage
-import os
 
 from flare.data.data_module import TestDataModule
 from flare.data.datasets import ContrastiveDataset
-from flare.utils.data import get_spec_featurizer, get_mol_featurizer, get_test_ms_dataset
+from flare.definitions import DATA_DIR, TEST_RESULTS_DIR
+from flare.utils.config import default_param_path, load_param_file
+from flare.utils.data import get_mol_featurizer, get_spec_featurizer, get_test_ms_dataset
 from flare.utils.models import get_model
 
-from flare.definitions import TEST_RESULTS_DIR
-import yaml
-from functools import partial
-# Suppress RDKit warnings and errors
 lg = RDLogger.logger()
 lg.setLevel(RDLogger.CRITICAL)
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--param_pth", type=str, default="params_formSpec.yaml")
-parser.add_argument('--checkpoint_pth', type=str, default='')
-parser.add_argument('--checkpoint_choice', type=str, default='train', choices=['train', 'val'])
-parser.add_argument('--df_test_pth', type=str, help='result file name')
-parser.add_argument('--exp_dir', type=str)
-parser.add_argument('--candidates_pth', type=str)
-parser.add_argument('--external_test', action='store_true', help='whether the test set is external data without labels')
+parser.add_argument(
+    "--param_pth",
+    type=str,
+    default=None,
+    help="YAML hyperparameters (default: FLARE_PARAMS or repo params.yaml)",
+)
+parser.add_argument("--checkpoint_pth", type=str, default="")
+parser.add_argument("--checkpoint_choice", type=str, default="train", choices=["train", "val"])
+parser.add_argument("--df_test_pth", type=str, help="result file name under experiment_dir")
+parser.add_argument("--exp_dir", type=str, help="experiment directory (overrides auto-detect)")
+parser.add_argument("--candidates_pth", type=str, help="override candidates JSON path")
+parser.add_argument(
+    "--external_test",
+    action="store_true",
+    help="external data without ground-truth labels in the candidate list",
+)
 
-def main(params):
-    # Seed everything
-    pl.seed_everything(params['seed'])
-        
-    # Init paths to data files
-    if params['debug']:
 
-        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')
+def main(params):
+    pl.seed_everything(params["seed"])
 
-    # Load dataset
-    spec_featurizer = get_spec_featurizer(params['spectra_view'], params)
+    if params.get("debug"):
+        dbg = os.environ.get("FLARE_DEBUG_DATASET")
+        if dbg:
+            params["dataset_pth"] = dbg
+        else:
+            sample_tsv = DATA_DIR / "sample" / "data.tsv"
+            if sample_tsv.is_file():
+                params["dataset_pth"] = str(sample_tsv)
+        params["split_pth"] = None
+        params["df_test_path"] = os.path.join(params["experiment_dir"], "debug_result.pkl")
 
-    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)
+    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
+    )
 
-    # Init data module
-    collate_fn = partial(ContrastiveDataset.collate_fn, spec_enc=params['spec_enc'], spectra_view=params['spectra_view'], stage=Stage.TEST)
+    collate_fn = partial(
+        ContrastiveDataset.collate_fn,
+        spec_enc=params["spec_enc"],
+        spectra_view=params["spectra_view"],
+        stage=Stage.TEST,
+    )
     data_module = TestDataModule(
         dataset=dataset,
         collate_fn=collate_fn,
-        split_pth=params['split_pth'],
-        batch_size=params['batch_size'],
-        num_workers=params['num_workers']
+        split_pth=params["split_pth"],
+        batch_size=params["batch_size"],
+        num_workers=params["num_workers"],
     )
 
-    model = get_model(params['model'], params)
-    model.df_test_path = params['df_test_path']
-    model.external_test = params['external_test']
-    
-    # Init trainer
+    model = get_model(params["model"], params)
+    model.df_test_path = params["df_test_path"]
+    model.external_test = params["external_test"]
+
     trainer = Trainer(
-        accelerator=params['accelerator'],
-        devices=params['devices'],
-        default_root_dir=params['experiment_dir']
+        accelerator=params["accelerator"],
+        devices=params["devices"],
+        default_root_dir=params["experiment_dir"],
     )
 
-    # Prepare data module to test
     data_module.prepare_data()
     data_module.setup(stage="test")
-        
-    # Test
     trainer.test(model, datamodule=data_module)
 
 
 if __name__ == "__main__":
     args = parser.parse_args([] if "__file__" not in globals() else None)
 
-    # Load
-    with open(args.param_pth) as f:
-        params = yaml.load(f, Loader=yaml.FullLoader)
-    
-    # Experiment directory
+    param_path = args.param_pth or str(default_param_path())
+    params = load_param_file(param_path)
+
+    exp_dir = None
     if args.exp_dir:
         exp_dir = args.exp_dir
     else:
-        run_name = params['run_name']
-        for exp in os.listdir(TEST_RESULTS_DIR): # find exp dir with matching run_name
-            if exp.endswith("_"+run_name):
-                exp_dir = str(TEST_RESULTS_DIR / exp)
-                break
-    if not exp_dir:
-        now = datetime.datetime.now().strftime("%Y%m%d")
-        exp_dir = str(TEST_RESULTS_DIR / f"{now}_{params['run_name']}")
-        os.makedirs(exp_dir, exist_ok=True)
-    print("EXPERIMENT directory: ",exp_dir)
-    params['experiment_dir'] = exp_dir
-    
-    # Checkpoint path
+        run_name = params["run_name"]
+        if TEST_RESULTS_DIR.is_dir():
+            for exp in sorted(os.listdir(TEST_RESULTS_DIR), reverse=True):
+                if exp.endswith("_" + run_name):
+                    exp_dir = str(TEST_RESULTS_DIR / exp)
+                    break
+        if exp_dir is None:
+            today_str = datetime.datetime.now().strftime("%Y%m%d")
+            exp_dir = str(TEST_RESULTS_DIR / f"{today_str}_{run_name}")
+            os.makedirs(exp_dir, exist_ok=True)
+    params["experiment_dir"] = exp_dir
+
     if args.checkpoint_pth:
-        params['checkpoint_pth'] = args.checkpoint_pth
-    
-    if not params['checkpoint_pth']:
-        print("No checkpoint provided. Using the checkpoint in the experiment directory")
-        for f in os.listdir(exp_dir):
+        params["checkpoint_pth"] = args.checkpoint_pth
+
+    if not params.get("checkpoint_pth"):
+        print("No checkpoint in params; searching experiment_dir for a .ckpt file")
+        for f in sorted(os.listdir(exp_dir)):
             if f.endswith("ckpt") and f.startswith("epoch") and args.checkpoint_choice in f:
-                checkpoint_path = os.path.join(exp_dir, f)
-                params['checkpoint_pth'] = checkpoint_path
+                params["checkpoint_pth"] = os.path.join(exp_dir, f)
                 break
-    assert(params['checkpoint_pth'] != '')
+    assert params.get("checkpoint_pth"), "No checkpoint found; pass --checkpoint_pth"
 
     if args.external_test:
-        params['external_test'] = True
+        params["external_test"] = True
     else:
-        params['external_test'] = False
-        
+        params["external_test"] = params.get("external_test", False)
+
     if args.candidates_pth:
-        params['candidates_pth'] = args.candidates_pth
+        params["candidates_pth"] = args.candidates_pth
+        from flare.utils.config import resolve_repo_paths
+
+        resolve_repo_paths(params)
+
     if args.df_test_pth:
-        params['df_test_path'] = os.path.join(exp_dir, args.df_test_pth)
-    if not params['df_test_path']:
-        params['df_test_path'] = os.path.join(exp_dir, f"result_{params['candidates_pth'].split('/')[-1].split('.')[0]}.pkl")
-        
+        params["df_test_path"] = os.path.join(exp_dir, args.df_test_pth)
+    if not params.get("df_test_path"):
+        cand = params.get("candidates_pth") or "candidates.json"
+        stem = Path(cand).stem
+        params["df_test_path"] = os.path.join(exp_dir, f"result_{stem}.pkl")
+
+    print("DF TEST PATH: ", params["df_test_path"])
+    print("EXP DIR: ", exp_dir)
+
     main(params)

flare/train.py

@@ -1,137 +1,128 @@
 import argparse
 import datetime
-
 import os
 import sys
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
-from rdkit import RDLogger
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from functools import partial
+
 import pytorch_lightning as pl
+import yaml
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks.early_stopping import EarlyStopping
-
+from rdkit import RDLogger
 
 from flare.data.data_module import ContrastiveDataModule
-
-from flare.definitions import TEST_RESULTS_DIR
-import yaml
 from flare.data.datasets import ContrastiveDataset
-from functools import partial
-
-from flare.utils.data import get_ms_dataset, get_spec_featurizer, get_mol_featurizer
+from flare.definitions import DATA_DIR, TEST_RESULTS_DIR
+from flare.utils.config import default_param_path, load_param_file
+from flare.utils.data import get_ms_dataset, get_mol_featurizer, get_spec_featurizer
 from flare.utils.models import get_model
-# Suppress RDKit warnings and errors
+
 lg = RDLogger.logger()
 lg.setLevel(RDLogger.CRITICAL)
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--param_pth", type=str, default="params_formSpec.yaml")
+parser.add_argument(
+    "--param_pth",
+    type=str,
+    default=None,
+    help="YAML hyperparameters (default: FLARE_PARAMS env or repo params.yaml)",
+)
+
 
 def main(params):
-    # Seed everything
-    pl.seed_everything(params['seed'])
-        
-    # Init paths to data files
-    if params['debug']:
-        params['dataset_pth'] = "/data/yzhouc01/MVP/data/sample/data.tsv"
-        params['candidates_pth'] =None
-        params['split_pth']=None
-
-    # Load dataset
-    spec_featurizer = get_spec_featurizer(params['spectra_view'], params)
-    mol_featurizer = get_mol_featurizer(params['molecule_view'], params)
-    dataset = get_ms_dataset(params['spectra_view'], params['molecule_view'], spec_featurizer, mol_featurizer, params)
-    
-    # Init data module
-    collate_fn = partial(ContrastiveDataset.collate_fn, spec_enc=params['spec_enc'], spectra_view=params['spectra_view'])
+    pl.seed_everything(params["seed"])
+
+    if params.get("debug"):
+        dbg = os.environ.get("FLARE_DEBUG_DATASET")
+        if dbg:
+            params["dataset_pth"] = dbg
+        else:
+            sample_tsv = DATA_DIR / "sample" / "data.tsv"
+            if sample_tsv.is_file():
+                params["dataset_pth"] = str(sample_tsv)
+        params["candidates_pth"] = None
+        params["split_pth"] = None
+
+    spec_featurizer = get_spec_featurizer(params["spectra_view"], params)
+    mol_featurizer = get_mol_featurizer(params["molecule_view"], params)
+    dataset = get_ms_dataset(
+        params["spectra_view"], params["molecule_view"], spec_featurizer, mol_featurizer, params
+    )
+
+    collate_fn = partial(
+        ContrastiveDataset.collate_fn, spec_enc=params["spec_enc"], spectra_view=params["spectra_view"]
+    )
     data_module = ContrastiveDataModule(
         dataset=dataset,
         collate_fn=collate_fn,
-        split_pth=params['split_pth'],
-        batch_size=params['batch_size'],
-        num_workers=params['num_workers'],
+        split_pth=params["split_pth"],
+        batch_size=params["batch_size"],
+        num_workers=params["num_workers"],
     )
 
-    model = get_model(params['model'], params)
-
-    # Init logger
-    if params['no_wandb']:
-        logger = None
-    else:
-        logger = pl.loggers.WandbLogger(
-            save_dir=params['experiment_dir'],
-            dir=params['experiment_dir'],
-            log_dir=params['experiment_dir'],
-            name=params['run_name'],
-            project=params['project_name'],
-            log_model=False,
-            config=model.hparams
-        )
+    model = get_model(params["model"], params)
+
+    tb_logger = pl.loggers.TensorBoardLogger(
+        save_dir=params["experiment_dir"],
+        name="",
+        version="",
+    )
+    logger = tb_logger
 
-    # Init callbacks for checkpointing and early stopping
-    callbacks = [pl.callbacks.ModelCheckpoint(save_last=False) ]
+    callbacks = [pl.callbacks.ModelCheckpoint(save_last=False)]
     for i, monitor in enumerate(model.get_checkpoint_monitors()):
-        monitor_name = monitor['monitor']
+        monitor_name = monitor["monitor"]
         checkpoint = pl.callbacks.ModelCheckpoint(
             monitor=monitor_name,
             save_top_k=1,
-            mode=monitor['mode'],
-            dirpath=params['experiment_dir'],
-            filename=f'{{epoch}}-{{{monitor_name}:.2f}}',
-            # filename='{epoch}-{val_loss:.2f}-{train_loss:.2f}',
+            mode=monitor["mode"],
+            dirpath=params["experiment_dir"],
+            filename=f"{{epoch}}-{{{monitor_name}:.2f}}",
             auto_insert_metric_name=True,
-            # save_last=(i == 0)
         )
         callbacks.append(checkpoint)
-        if monitor.get('early_stopping', False):
+        if monitor.get("early_stopping", False):
             early_stopping = EarlyStopping(
                 monitor=monitor_name,
-                mode=monitor['mode'],
+                mode=monitor["mode"],
                 verbose=True,
-                patience=params['early_stopping_patience'],
+                patience=params["early_stopping_patience"],
             )
             callbacks.append(early_stopping)
 
-    # Init trainer
     trainer = Trainer(
-        accelerator=params['accelerator'],
-        devices=params['devices'],
-        max_epochs=params['max_epochs'],
+        accelerator=params["accelerator"],
+        devices=params["devices"],
+        max_epochs=params["max_epochs"],
         logger=logger,
-        log_every_n_steps=params['log_every_n_steps'],
-        val_check_interval=params['val_check_interval'],
+        log_every_n_steps=params["log_every_n_steps"],
+        val_check_interval=params["val_check_interval"],
         callbacks=callbacks,
-        default_root_dir=params['experiment_dir'],
+        default_root_dir=params["experiment_dir"],
     )
 
-    # Prepare data module to validate or test before training
     data_module.prepare_data()
     data_module.setup()
 
-
-    # Validate before training
     trainer.validate(model, datamodule=data_module)
-
-    # Train
     trainer.fit(model, datamodule=data_module)
-        
 
 
 if __name__ == "__main__":
     args = parser.parse_args([] if "__file__" not in globals() else None)
 
-    # Get current time
+    param_path = args.param_pth or str(default_param_path())
+    params = load_param_file(param_path)
+
     now = datetime.datetime.now()
     now_formatted = now.strftime("%Y%m%d")
-
-    # Load
-    with open(args.param_pth) as f:
-        params = yaml.load(f, Loader=yaml.FullLoader)
-
     experiment_dir = str(TEST_RESULTS_DIR / f"{now_formatted}_{params['run_name']}")
-    params['experiment_dir'] = experiment_dir
+    params["experiment_dir"] = experiment_dir
 
-    if not params['df_test_path']:
-        params['df_test_path'] = os.path.join(experiment_dir, "result.pkl")
+    if not params.get("df_test_path"):
+        params["df_test_path"] = os.path.join(experiment_dir, "result.pkl")
 
     main(params)

flare/tune.py

@@ -1,4 +1,5 @@
 import argparse
+import copy
 import datetime
 import os
 import sys
@@ -20,6 +21,7 @@ from flare.data.datasets import ContrastiveDataset
 from flare.utils.data import get_ms_dataset, get_spec_featurizer, get_mol_featurizer
 from flare.utils.models import get_model
 from flare.definitions import TEST_RESULTS_DIR
+from flare.utils.config import default_param_path, load_param_file
 from functools import partial
 from rdkit import RDLogger
 from massspecgym.models.base import Stage
@@ -29,7 +31,12 @@ lg = RDLogger.logger()
 lg.setLevel(RDLogger.CRITICAL)
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--param_pth", type=str, default="params_formSpec.yaml")
+parser.add_argument(
+    "--param_pth",
+    type=str,
+    default=None,
+    help="Base YAML (default: FLARE_PARAMS or repo params.yaml)",
+)
 parser.add_argument("--n_trials", type=int, default=20)
 
 class EpochLossTracker(Callback):
@@ -112,7 +119,7 @@ def save_trial_result(base_dir, trial, params, duration):
 
 def objective(trial: optuna.Trial, base_params, trial_times, base_dir, total_trials):
     start_time = time.time()
-    params = base_params.copy()
+    params = copy.deepcopy(base_params)
 
     try:
         # Training-related params
@@ -160,8 +167,6 @@ def objective(trial: optuna.Trial, base_params, trial_times, base_dir, total_tri
             ContrastiveDataset.collate_fn,
             spec_enc=params["spec_enc"],
             spectra_view=params["spectra_view"],
-            mask_peak_ratio=params["mask_peak_ratio"],
-            aug_cands=params["aug_cands"],
         )
 
         data_module = ContrastiveDataModule(
@@ -226,12 +231,11 @@ def objective(trial: optuna.Trial, base_params, trial_times, base_dir, total_tri
 
 
 def main(args):
-    with open(args.param_pth) as f:
-        params = yaml.load(f, Loader=yaml.FullLoader)
+    param_path = args.param_pth or str(default_param_path())
+    params = load_param_file(param_path)
 
-    # now = datetime.datetime.now().strftime("%Y%m%d")
-    # base_dir = str(TEST_RESULTS_DIR / f"{now}_{params['run_name']}_optuna")
-    base_dir = "../experiments/20250916_simple_model_optuna"
+    now = datetime.datetime.now().strftime("%Y%m%d")
+    base_dir = str(TEST_RESULTS_DIR / f"{now}_{params['run_name']}_optuna")
     os.makedirs(base_dir, exist_ok=True)
     params["experiment_dir"] = base_dir
 

flare/utils/__init__.py

@@ -1,3 +1 @@
-import sys
-sys.path.insert(0, "/data/yzhouc01/MassSpecGym")
 from massspecgym.utils import *

flare/utils/config.py

@@ -0,0 +1,63 @@
+"""Load YAML hyperparameters and resolve filesystem paths relative to the repository root."""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from flare.definitions import REPO_DIR
+
+# Keys that may hold filesystem paths (relative paths are resolved against REPO_DIR).
+_PATH_KEYS = frozenset(
+    {
+        "dataset_pth",
+        "candidates_pth",
+        "subformula_dir_pth",
+        "split_pth",
+        "checkpoint_pth",
+        "df_test_path",
+        "formula_to_smiles_pth",
+    }
+)
+
+
+def resolve_repo_paths(params: dict[str, Any]) -> None:
+    """In-place: turn repo-relative path strings into absolute paths."""
+    root = REPO_DIR
+    for key in _PATH_KEYS:
+        val = params.get(key)
+        if not val or not isinstance(val, str):
+            continue
+        p = Path(val)
+        if not p.is_absolute():
+            params[key] = str((root / p).resolve())
+        else:
+            params[key] = str(p.resolve())
+
+
+def load_param_file(path: str | Path) -> dict[str, Any]:
+    """Load a YAML parameter file and resolve path fields."""
+    p = Path(path)
+    if not p.is_file():
+        raise FileNotFoundError(f"Parameter file not found: {p}")
+    with open(p, encoding="utf-8") as f:
+        params = yaml.load(f, Loader=yaml.FullLoader)
+    if params is None:
+        params = {}
+    if not isinstance(params, dict):
+        raise TypeError(f"Expected mapping at top level of {p}, got {type(params)}")
+    resolve_repo_paths(params)
+    return params
+
+
+def default_param_path() -> Path:
+    """Path to the default params file (overridable with FLARE_PARAMS)."""
+    override = os.environ.get("FLARE_PARAMS")
+    if override:
+        return Path(override).expanduser()
+    env_root = os.environ.get("FLARE_REPO_ROOT")
+    if env_root:
+        return Path(env_root).expanduser() / "params.yaml"
+    return REPO_DIR / "params.yaml"

flare/utils/data.py

@@ -21,6 +21,11 @@ class Subformula_Loader:
         self.dir_path = dir_path
         self.use_prec_mz = use_prec_mz
         self.formula_source = formula_source
+        if str(formula_source).startswith('magma'):
+            raise ValueError(
+                "MAGMA formula sources are not supported in this release (see archive/magma/). "
+                "Use 'default' (MIST-style JSON) or 'sirius'."
+            )
         if spectra_view == 'SpecFormula':
             self.load = self.load_subformula_data
         elif spectra_view == "SpecFormulaMz":
@@ -63,77 +68,6 @@ class Subformula_Loader:
         except:
             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 {'formulas': formulas, 'formula_mzs': mzs, 'formula_intensities': intensities}
-
     def load_sirius_data(self, data):
         try:
 
@@ -167,8 +101,6 @@ class Subformula_Loader:
                 data = json.load(f)
             if self.formula_source == 'sirius':
                 return self.load_sirius_data(data)
-            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)
 
@@ -263,17 +195,18 @@ def get_test_ms_dataset(spectra_view: T.Union[str, T.List[str]],
         else: views.extend(v)
     views = frozenset(views)
 
-    dataset_params = {'spectra_view': spectra_view, 'pth': params['dataset_pth'], 'spec_transform': spectra_featurizer, 'mol_transform': mol_featurizer, "candidates_pth": params['candidates_pth']}
+    dataset_params = {
+        'spectra_view': spectra_view,
+        'pth': params['dataset_pth'],
+        'spec_transform': spectra_featurizer,
+        'mol_transform': mol_featurizer,
+        "candidates_pth": params.get('candidates_pth'),
+        "formula_to_smiles_pth": params.get('formula_to_smiles_pth'),
+        "external_test": params.get('external_test', False)
+    }
     if "SpecFormula" in views or "SpecFormulaMz" in views:
-        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']:
-    #     dataset_params.update({'cons_spec_dir_pth': params['cons_spec_dir_pth']})
-    # if 'use_NL_spec' in params and params['use_NL_spec']:
-    #     dataset_params.update({'NL_spec_dir_pth': params['NL_spec_dir_pth']})
-    # if params['pred_fp'] or params['use_fp']:
-    #     dataset_params.update({'fp_dir_pth': '', 'fp_size': params['fp_size'], 'fp_radius': params['fp_radius']})
+        dataset_params.update({'subformula_dir_pth': params['subformula_dir_pth'], 'formula_source': params['formula_source']})
+        use_formulas = True
 
     return jestr_datasets.ExpandedRetrievalDataset(use_formulas=use_formulas, **dataset_params)
     

flare/utils/loss.py

@@ -43,40 +43,6 @@ def contrastive_loss(v1, v2, tau=1.0) -> torch.Tensor:
         
         return Lv1_v2 + Lv2_v1 , torch.mean(numerator), torch.mean(Lv1_v2_denom+Lv2_v1_denom)
 
-def cand_spec_sim_loss(spec_enc, cand_enc):
-        cand_enc = torch.transpose(cand_enc, 0, 1) # C x B x d
-        spec_enc = spec_enc.unsqueeze(0) # 1 x B x d
-
-        sim = nn.functional.cosine_similarity(spec_enc, cand_enc, dim=2)
-        loss = torch.mean(sim)
-
-        return loss
-
-class cons_spec_loss:
-        def __init__(self, loss_type) -> None:
-                self.loss_compute = {'cosine': self.cos_loss,
-                                     'l2':torch.nn.MSELoss()}[loss_type]
-        def __call__(self,cons_spec, ind_spec):
-                return self.loss_compute(cons_spec, ind_spec)
-        
-        def cos_loss(self, cons_spec, ind_spec):
-                sim = nn.functional.cosine_similarity(cons_spec, ind_spec)
-                loss = 1-torch.mean(sim) 
-                return loss
-
-class fp_loss:
-        def __init__(self, loss_type) -> None:
-                self.loss_compute = {'cosine': self.fp_loss_cos,
-                                        'bce': nn.BCELoss()}[loss_type]
-        
-        def __call__(self, predicted_fp, target_fp):
-                return self.loss_compute(predicted_fp, target_fp)
-        
-        def fp_loss_cos(self, predicted_fp, target_fp):
-                sim = nn.functional.cosine_similarity(predicted_fp, target_fp)
-                return 1 - torch.mean(sim)
-
-
 # ---------- Utility ----------
 def _safe_divide(num, denom, eps=1e-8):
     return num / (denom + eps)

flare/utils/models.py

@@ -1,6 +1,5 @@
 from flare.models.spec_encoder import SpecEncMLP_BIN, SpecFormulaEncMLP, SpecFormulaTransformer,SpecFormula_mz_Encoder, SpecMzIntTokenTransformer
 from flare.models.mol_encoder import MolEnc
-from flare.models.encoders import MLP
 from flare.models.contrastive import ContrastiveModel, CrossAttenContrastive, FilipContrastive, FilipGlobalContrastive
 
 def get_spec_encoder(spec_enc:str, args):
@@ -13,12 +12,6 @@ def get_spec_encoder(spec_enc:str, args):
 def get_mol_encoder(mol_enc: str, args):
     return {'GNN': MolEnc}[mol_enc](args, in_dim=78)
 
-def get_fp_pred_model(args):
-    return MLP(in_dim=args.final_embedding_dim, hidden_dims=[args.fp_size], final_activation='sigmoid', dropout=args.fp_dropout)
-
-def get_fp_enc_model(args):
-    return MLP(in_dim=args.fp_size, hidden_dims=[args.final_embedding_dim,args.final_embedding_dim*2,args.final_embedding_dim,], final_activation=None, dropout=0.0)
-
 def get_model(model:str,
               params):
     

flare/utils/mol_search.py

@@ -311,57 +311,7 @@ class SpectraMoleculeRetriever:
 
 
 if __name__ == "__main__":
-    import sys
-    sys.path.insert(0, "/data/yzhouc01/FILIP-MS")
-
-    from flare.utils.data import get_spec_featurizer, get_mol_featurizer
-    from flare.utils.models import get_model
-    from flare.utils.mol_search import SpectraMoleculeRetriever
-    from flare.utils.general import filip_similarity_single
-    import yaml
-
-    metadata = {
-        "class": {
-            "lipid": ["mol1", "mol2"],
-            "peptide": ["mol3"]
-        },
-        "pathway": {
-            "beta-oxidation": ["mol1"],
-            "glycolysis": ["mol2", "mol3"]
-        }
-    }
-
-    smiles_dict = {
-        "mol1": "CCO",
-        "mol2": "CCN",
-        "mol3": "CCC"
-    }
-
-    # Load model and data
-    param_pth = '/data/yzhouc01/cancer/flare.yaml'
-    with open(param_pth) as f:
-        params = yaml.load(f, Loader=yaml.FullLoader)
-
-    spec_featurizer = get_spec_featurizer(params['spectra_view'], params)
-    mol_featurizer = get_mol_featurizer(params['molecule_view'], params)
-
-
-    # load model
-    checkpoint_pth = "/data/yzhouc01/FILIP-MS/experiments/20250930_optimized_flare_42/epoch=1959-train_loss=0.08.ckpt"
-    params['checkpoint_pth'] = checkpoint_pth
-    model = get_model(params['model'], params)
-
-    specMolRetriever = SpectraMoleculeRetriever(
-        molecule_encoder=model.mol_enc_model,
-        spectra_encoder=model.spec_enc_model,
-        fine_similarity_fn=filip_similarity_single,
-        smiles_preprocess=mol_featurizer
+    raise SystemExit(
+        "SpectraMoleculeRetriever is a library class; configure paths via FLARE_PARAMS / "
+        "FLARE_CHECKPOINT and import it from application code (see README)."
     )
-
-    specMolRetriever.build_database(smiles_dict, metadata=metadata, cache_nodes=True)
-
-    # Filter search to molecules in a specific pathway
-    # results = specMolRetriever.search(spectrum, subset={"pathway": "beta-oxidation"})
-
-    # for mol_id, score in results[:10]:
-    #     print(f"{mol_id}: {score:.3f}")