model code

mvp/data/data_module.py

@@ -0,0 +1,84 @@
+from torch.utils.data.dataloader import DataLoader
+from massspecgym.data.data_module import MassSpecDataModule
+from mvp.data.datasets import ContrastiveDataset
+from functools import partial
+from massspecgym.models.base import Stage
+
+class TestDataModule(MassSpecDataModule):
+    def __init__(
+            self,
+            collate_fn,
+            **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.collate_fn = collate_fn
+
+    def prepare_data(self):
+        pass
+    
+    def setup(self, stage=None):
+        if stage == "test":
+            self.test_dataset = self.dataset
+        else:
+            raise Exception("Data module supports test set only")
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            persistent_workers=self.persistent_workers,
+            drop_last=False,
+            collate_fn=self.collate_fn,
+        )
+
+    def train_dataloader(self):
+        return None
+    
+    def val_dataset(self):
+        return None
+
+class ContrastiveDataModule(MassSpecDataModule):
+    def __init__(
+            self,
+            collate_fn,
+            **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.collate_fn = collate_fn
+        self.regularization_flag = False
+             
+    def train_dataloader(self):
+        self.train_contrastive_dataset = ContrastiveDataset(self.train_dataset)
+
+        return DataLoader(self.train_contrastive_dataset,
+                          batch_size=self.batch_size,
+                          shuffle=True,
+                          num_workers=self.num_workers,
+                          persistent_workers=self.persistent_workers,
+                          drop_last=False,
+                          collate_fn=partial(self.collate_fn, stage=Stage.TRAIN),
+                          )
+
+    def val_dataloader(self):
+        self.val_contrastive_dataset = ContrastiveDataset(self.val_dataset)
+
+        return DataLoader(self.val_contrastive_dataset,
+                          batch_size=self.batch_size,
+                          shuffle=False,
+                          num_workers=self.num_workers,
+                          persistent_workers=self.persistent_workers,
+                          drop_last=False,
+                          collate_fn=partial(self.collate_fn, stage=Stage.VAL))
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            persistent_workers=self.persistent_workers,
+            drop_last=False,
+            collate_fn=self.dataset.collate_fn,
+        )

mvp/data/datasets.py

@@ -0,0 +1,403 @@
+import pandas as pd
+import json
+import typing as T
+import numpy as np
+import torch
+import massspecgym.utils as utils
+from pathlib import Path
+from torch.utils.data.dataset import Dataset
+from torch.utils.data.dataloader import default_collate
+import dgl
+from collections import defaultdict
+from massspecgym.data.transforms import SpecTransform, MolTransform, MolToInChIKey
+from massspecgym.data.datasets import MassSpecDataset
+import mvp.utils.data as data_utils
+from torch.nn.utils.rnn import pad_sequence
+from massspecgym.models.base import Stage
+import pickle
+import math
+import itertools
+from rdkit.Chem import AllChem
+from rdkit import Chem
+class JESTR1_MassSpecDataset(MassSpecDataset):
+    def __init__(
+        self,
+        spectra_view: str,
+        fp_dir_pth: str = None,
+        cons_spec_dir_pth: str = None,
+        NL_spec_dir_pth: str = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        self.use_fp = False
+        self.use_cons_spec = False
+        self.use_NL_spec = False
+        self.spectra_view = spectra_view
+
+        # load fingerprints
+        self._load_fp(fp_dir_pth)
+
+        # load consensus
+        self._load_cons_spec(cons_spec_dir_pth)
+
+        # load NL specs
+        self._load_NL_spec(NL_spec_dir_pth)
+
+    def _load_fp(self, fp_dir_pth):
+        if fp_dir_pth is not None:
+            self.use_fp = True
+            if fp_dir_pth:
+                with open(fp_dir_pth, 'rb') as f:
+                    self.smiles_to_fp = pickle.load(f)
+            else:
+                self.smiles_to_fp = {}
+    
+    def _load_cons_spec(self, cons_spec_dir_pth):
+        if cons_spec_dir_pth is not None:
+            self.use_cons_spec = True
+            with open(cons_spec_dir_pth, 'rb') as f:
+                cons_specs = pickle.load(f)
+
+            # Convert spectra to matchms spectra
+            matchMS_preparer = data_utils.PrepMatchMS(self.spectra_view)
+            spectra = cons_specs.apply(matchMS_preparer.prepare,axis=1)
+
+            self.cons_specs = dict(zip(cons_specs['smiles'].tolist(), spectra))
+
+    def _load_NL_spec(self, NL_spec_dir_pth):
+        if NL_spec_dir_pth is not None:
+            self.use_NL_spec = True
+            with open(NL_spec_dir_pth, 'rb') as f:
+                NL_specs = pickle.load(f)
+
+            # Convert spectra to matchms spectra
+            matchMS_preparer = data_utils.PrepMatchMS(self.spectra_view)
+            self.NL_specs = NL_specs.apply(matchMS_preparer.prepare,axis=1)
+
+
+    def __getitem__(self, i, transform_spec: bool = True, transform_mol: bool = True):
+
+        spec = self.spectra[i]
+        metadata = self.metadata.iloc[i]
+        mol = metadata["smiles"]
+
+        # 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)
+        else:
+            item["spec"] = spec
+
+        if self.return_mol_freq:
+            item["mol_freq"] = metadata["mol_freq"]
+
+        if self.return_identifier:
+            item["identifier"] = metadata["identifier"]
+
+        if self.use_fp and self.smiles_to_fp:
+            item['fp'] = torch.Tensor(self.smiles_to_fp[mol].ToList())
+        
+        if self.use_cons_spec:
+            item['cons_spec'] = self.spec_transform[self.spectra_view](self.cons_specs[mol])
+
+        if self.use_NL_spec:
+            item['NL_spec'] = self.spec_transform[self.spectra_view](self.NL_specs[i])
+
+        # Apply all transformations to the molecule
+        if transform_mol and self.mol_transform:
+            if isinstance(self.mol_transform, dict):
+                for key, transform in self.mol_transform.items():
+                    item[key] = transform(mol) if transform is not None else mol
+            else:
+                item["mol"] = self.mol_transform(mol)
+        else:
+            item["mol"] = mol
+        return item
+
+class MassSpecDataset_PeakFormulas(JESTR1_MassSpecDataset):
+    def __init__(
+        self,
+        spectra_view: str,
+        spec_transform: T.Optional[T.Union[SpecTransform, T.Dict[str, SpecTransform]]],
+        mol_transform: T.Optional[T.Union[MolTransform, T.Dict[str, MolTransform]]],
+        pth: T.Optional[Path],
+        subformula_dir_pth: str,
+        fp_dir_pth: str = None,
+        NL_spec_dir_pth: str = None,
+        cons_spec_dir_pth: str = None,
+        return_mol_freq: bool = False,
+        return_identifier: bool = True,
+        dtype: T.Type = torch.float32
+    ):
+        """
+        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
+
+        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")
+
+        # 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)
+
+        # 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())))
+        
+        
+        # load fingerprints
+        self._load_fp(fp_dir_pth)
+
+        # load consensus spectra
+        self._load_cons_spec(cons_spec_dir_pth)
+
+        # load NL specs
+        self._load_NL_spec(NL_spec_dir_pth)
+
+        self.metadata = self.metadata[self.metadata['identifier'].isin(id_to_spec)]
+        formula_df = pd.DataFrame.from_dict(id_to_spec, orient='index').reset_index().rename(columns={'index': 'identifier'})
+        self.metadata = self.metadata.merge(formula_df, on='identifier')
+
+        # create matchms spectra
+        matchMS_preparer = data_utils.PrepMatchMS(spectra_view=spectra_view)
+        self.spectra = self.metadata.apply(matchMS_preparer.prepare,axis=1)
+                
+        if self.return_mol_freq:
+            if "inchikey" not in self.metadata.columns:
+                self.metadata["inchikey"] = self.metadata["smiles"].apply(utils.smiles_to_inchi_key)
+            self.metadata["mol_freq"] = self.metadata.groupby("inchikey")["inchikey"].transform("count")
+
+        self.return_identifier = return_identifier
+        self.dtype = dtype
+    
+    def __getitem__(self, i, transform_spec: bool = True, transform_mol: bool = True):
+        item = super().__getitem__(i, transform_spec, transform_mol = False)
+        mol = item['mol'] #smiles
+
+        # transform mol
+        if transform_mol:
+            if isinstance(self.mol_transform, dict):
+                for key, transform in self.mol_transform.items():
+                    item[key] = transform(mol) if transform is not None else mol
+            else:
+                item["mol"] = self.mol_transform(mol)
+
+        return item
+
+class ContrastiveDataset(Dataset):
+    def __init__(
+        self,
+        spec_mol_data,
+    ):
+        super().__init__()
+    
+        indices = spec_mol_data.indices
+        self.spec_mol_data = spec_mol_data
+        self.smiles_to_specmol_ids = spec_mol_data.dataset.metadata.loc[indices].groupby('smiles').indices
+        self.smiles_to_spec_couter = defaultdict(int)
+        self.smiles_list = list(self.smiles_to_specmol_ids.keys())
+
+    def __len__(self) -> int:
+        return len(self.smiles_list)
+    
+    def __getitem__(self, i:int) -> dict:
+        mol = self.smiles_list[i]
+
+        # select spectrum (iterate through list of spectra)
+        specmol_ids = self.smiles_to_specmol_ids[mol]
+        counter = self.smiles_to_spec_couter[mol]
+        specmol_id = specmol_ids[counter % len(specmol_ids)]
+
+        item = self.spec_mol_data.__getitem__(specmol_id)
+        self.smiles_to_spec_couter[mol] = counter+1
+        # item['smiles'] = mol
+        # item['spec_id'] = specmol_id
+        return item
+
+    @staticmethod
+    def collate_fn(batch: T.Iterable[dict], spec_enc: str, spectra_view: str, stage=None, mask_peak_ratio: float = 0.0, aug_cands: bool = False) -> dict:
+        mol_key = 'cand' if stage == Stage.TEST else 'mol'
+        non_standard_collate = ['mol', 'cand', 'aug_cands', 'cons_spec', 'aug_cands_fp', 'NL_spec']
+        require_pad = False
+        if 'Formula' in spectra_view or 'Tokens' in spectra_view:
+            require_pad = True
+            padding_value=-5 if spec_enc in ('Transformer_Formula', 'Formula_BinnedSpec', 'Transformer_MzInt') else 0
+            non_standard_collate.append(spectra_view)
+        else:
+            non_standard_collate.remove('cons_spec')
+            non_standard_collate.remove('NL_spec')
+
+        collated_batch = {}
+        # standard collate
+        for k in batch[0].keys():
+            if k not in non_standard_collate:
+                collated_batch[k] = default_collate([item[k] for item in batch])
+                
+        # batch graphs
+        batch_mol = []
+        batch_mol_nodes= []
+
+        for item in batch:
+            batch_mol.append(item[mol_key])
+            batch_mol_nodes.append(item[mol_key].num_nodes())
+
+        collated_batch[mol_key] = dgl.batch(batch_mol)
+        collated_batch['mol_n_nodes'] = batch_mol_nodes
+        
+        # pad peaks/formulas
+        if require_pad:
+            peaks = []
+            n_peaks = []
+            for item in batch:
+                peaks.append(item[spectra_view])
+                n_peaks.append(len(item[spectra_view]))
+            collated_batch[spectra_view] = pad_sequence(peaks, batch_first=True, padding_value=padding_value)
+            collated_batch['n_peaks'] = n_peaks
+        
+            if 'cons_spec' in batch[0]:
+                peaks = []
+                n_peaks = []
+                for item in batch:
+                    peaks.append(item['cons_spec'])
+                    n_peaks.append(len(item['cons_spec']))
+                collated_batch['cons_spec'] = pad_sequence(peaks, batch_first=True, padding_value=padding_value)
+                collated_batch['cons_n_peaks'] = n_peaks
+
+            if 'NL_spec' in batch[0]:
+                peaks = []
+                n_peaks = []
+                for item in batch:
+                    peaks.append(item['NL_spec'])
+                    n_peaks.append(len(item['NL_spec']))
+                collated_batch['NL_spec'] = pad_sequence(peaks, batch_first=True, padding_value=padding_value)
+                collated_batch['NL_n_peaks'] = n_peaks
+
+
+        # mask peaks
+        if mask_peak_ratio > 0.0 and stage == Stage.TRAIN:
+            n_mask_peaks = [math.floor(n_peak* mask_peak_ratio) for n_peak in n_peaks]
+            mask_peak_idx = [np.random.choice(n_peak, n_mask, replace=False) for n_peak, n_mask in zip(n_peaks, n_mask_peaks)]
+            for i, peaks in enumerate(collated_batch[spectra_view]):
+                peaks[mask_peak_idx[i]] = -5.0
+        
+        # batch candidates
+        if aug_cands:
+            candidates = \
+                sum([item["aug_cands"] for item in batch], start=[])
+            collated_batch['aug_cands'] = dgl.batch(candidates)
+
+            if 'aug_cands_fp' in batch[0]:
+                cand_fp = [item['aug_cands_fp'] for item in batch]
+                collated_batch['aug_cands_fp'] = torch.flatten(torch.Tensor(cand_fp), end_dim=1)
+
+        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,
+                 external_test: bool = False,
+                **kwargs):
+        
+        self.external_test = external_test
+        
+        self.instance = MassSpecDataset_PeakFormulas(**kwargs, return_mol_freq=False) if use_formulas else JESTR1_MassSpecDataset(**kwargs, return_mol_freq=False)
+        # super().__init__(**kwargs)
+
+        if self.use_fp:
+            self.fpgen = AllChem.GetMorganGenerator(radius=fp_radius,fpSize=fp_size)
+
+        self.candidates_pth = candidates_pth
+        self.mol_label_transform = mol_label_transform
+        
+        # Read candidates_pth from json to dict: SMILES -> respective candidate SMILES
+        with open(self.candidates_pth, "r") as file:
+            candidates = json.load(file)
+
+        self.candidates = {}
+        for s, cand in candidates.items():
+            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 self.external_test or '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']
+        test_smiles = self.metadata[self.metadata['fold'] == "test"]['smiles'].tolist()
+        test_ms_id = self.metadata[self.metadata['fold'] == "test"]['identifier'].tolist()
+        
+        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]
+            if not self.external_test:
+                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")
+            else:
+                labels = [False] * len(candidates)
+
+            self.spec_cand.extend([(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]
+
+        item = self.instance.__getitem__(spec_i, transform_mol=False)
+        item['cand'] = self.mol_transform(cand_smiles)
+        item['cand_smiles'] = cand_smiles
+        item['label'] = label
+
+        if self.use_fp:
+            item['fp'] = torch.Tensor(self.fpgen.GetFingerprint(Chem.MolFromSmiles(cand_smiles)).ToList())
+
+        return item

mvp/data/transforms.py

@@ -0,0 +1,202 @@
+import numpy as np
+import torch
+import matchms
+from typing import Optional
+from rdkit.Chem import AllChem as Chem
+from mvp.definitions import CHEM_ELEMS_SMALL
+from massspecgym.data.transforms import MolTransform, SpecTransform, default_matchms_transforms
+from massspecgym.data.transforms import SpecBinner
+
+import dgllife.utils as chemutils
+import re
+
+class SpecBinnerLog(SpecTransform):
+    def __init__(
+        self,
+        max_mz: float = 1005,
+        bin_width: float = 1,
+    ) -> None:
+        self.max_mz = max_mz
+        self.bin_width = bin_width
+        if not (max_mz / bin_width).is_integer():
+            raise ValueError("`max_mz` must be divisible by `bin_width`.")
+        
+    def matchms_transforms(self, spec: matchms.Spectrum) -> matchms.Spectrum:
+        return default_matchms_transforms(spec, mz_to=self.max_mz, n_max_peaks=None)
+    
+    def matchms_to_torch(self, spec: matchms.Spectrum) -> torch.Tensor:
+        """
+        Bin the spectrum into a fixed number of bins.
+        """
+        binned_spec = self._bin_mass_spectrum(
+            mzs=spec.peaks.mz,
+            intensities=spec.peaks.intensities,
+            max_mz=self.max_mz,
+            bin_width=self.bin_width,
+        )
+        return torch.from_numpy(binned_spec).to(dtype=torch.float32)
+
+    def _bin_mass_spectrum(
+        self, mzs, intensities, max_mz, bin_width
+    ):
+
+        # Calculate the number of bins
+        num_bins = int(np.ceil(max_mz / bin_width))
+
+        # Calculate the bin indices for each mass
+        bin_indices = np.floor(mzs -1 / bin_width).astype(int)
+
+        # Filter out mzs that exceed max_mz
+        valid_indices = bin_indices[mzs <= max_mz]
+        valid_intensities = intensities[mzs <= max_mz]
+
+        # Clip bin indices to ensure they are within the valid range
+        valid_indices = np.clip(valid_indices, 0, num_bins - 1)
+
+        # Initialize an array to store the binned intensities
+        binned_intensities = np.zeros(num_bins)
+
+        # Use np.add.at to sum intensities in the appropriate bins
+        np.add.at(binned_intensities, valid_indices, valid_intensities)
+
+        binned_intensities = binned_intensities/np.max(binned_intensities) * 999
+
+        binned_intensities = np.log10(binned_intensities + 1) / 3
+
+        return binned_intensities 
+    
+class SpecFormulaFeaturizer(SpecTransform):
+    ''' Uses processed mz and intensities, excludes mz values, keep peaks with formulas only'''
+    def __init__(
+            self,
+            add_intensities: bool,
+            max_mz: float = 1005,
+            element_list: list = CHEM_ELEMS_SMALL,
+            formula_normalize_vector: Optional[np.array] = None
+    ) -> None:
+        self.max_mz = max_mz
+        self.elem_to_pos = {e: i for i, e in enumerate(element_list)}
+        self.add_intensities = add_intensities
+        if formula_normalize_vector is None:
+            formula_normalize_vector = np.ones(len(element_list))
+        self.formula_normalize_vector = formula_normalize_vector
+        self.CHEM_FORMULA_SIZE = "([A-Z][a-z]*)([0-9]*)"
+        
+    def matchms_transforms(self, spec: matchms.Spectrum):
+        return spec
+    
+    def matchms_to_torch(self, spec: matchms.Spectrum) -> torch.Tensor:
+        mzs = spec.peaks.mz
+        intensities = spec.peaks.intensities
+        formulas = spec.metadata['formulas'] # list of formulas
+
+        peak_idx = np.where(mzs <= self.max_mz)[0]
+        intensities = intensities[peak_idx]
+        formulas = formulas[peak_idx]
+
+        spec = self._featurize_formula(formulas)
+        spec = spec/self.formula_normalize_vector
+
+        if self.add_intensities:
+            spec = np.concatenate((spec, intensities.reshape(-1,1)), axis=1)
+        spec = spec.astype(np.float32)
+
+        return torch.from_numpy(spec)
+    
+    def _featurize_formula(self, formulas):
+        formula_vector = np.zeros((len(formulas), len(self.elem_to_pos)))
+        for i, f in enumerate(formulas):
+            try:
+                for (e, ct) in re.findall(self.CHEM_FORMULA_SIZE, f):
+                    ct = 1 if ct == "" else int(ct)
+                    try:
+                        formula_vector[i][self.elem_to_pos[e]]+=ct
+                    except:
+                            print(f"Couldn't vectorize {f}, element {e} not supported")
+                            continue
+            except:
+                print(f"Couldn't vectorize {f}, formula not supported")
+                continue
+        return formula_vector
+
+class MolToGraph(MolTransform):
+    def __init__ (self, atom_feature: str = "full", bond_feature: str = "full", element_list: list = CHEM_ELEMS_SMALL):
+        self.atom_feature = atom_feature
+        self.bond_feature = bond_feature
+        self.node_featurizer = self._get_atom_featurizer(element_list=element_list) 
+        self.edge_featurizer = self._get_bond_featurizer()
+    
+    def from_smiles(self, mol:str):
+        mol = Chem.MolFromSmiles(mol)
+        g = chemutils.mol_to_bigraph(mol, node_featurizer=self.node_featurizer, edge_featurizer=self.edge_featurizer, add_self_loop = True,
+                             num_virtual_nodes = 0, canonical_atom_order=False)
+
+        # atom_ids = [atom.GetIdx() for atom in mol.GetAtoms()] # added for visualization
+        # g.ndata['atom_id'] = torch.tensor(atom_ids, dtype=torch.long)
+
+        return g
+
+    def _get_atom_featurizer(self, element_list) -> dict:
+        feature_mode = self.atom_feature
+        atom_mass_fun = chemutils.ConcatFeaturizer(
+            [chemutils.atom_mass]
+        )
+        def atom_bond_type_one_hot(atom):
+            bs = atom.GetBonds()
+            bt = np.array([chemutils.bond_type_one_hot(b) for b in bs])
+            return [any(bt[:, i]) for i in range(bt.shape[1])]
+
+        def atom_type_one_hot(atom):
+            return chemutils.atom_type_one_hot(
+                atom, allowable_set = element_list, encode_unknown = True
+            )
+        
+        if feature_mode == 'light':
+            atom_featurizer_funs = chemutils.ConcatFeaturizer([
+                chemutils.atom_mass,
+                atom_type_one_hot
+            ])
+        elif feature_mode == 'full':
+            atom_featurizer_funs = chemutils.ConcatFeaturizer([
+                chemutils.atom_mass,
+                atom_type_one_hot, 
+                atom_bond_type_one_hot,
+                chemutils.atom_degree_one_hot, 
+                chemutils.atom_total_degree_one_hot,
+                chemutils.atom_explicit_valence_one_hot,
+                chemutils.atom_implicit_valence_one_hot,
+                chemutils.atom_hybridization_one_hot,
+                chemutils.atom_total_num_H_one_hot,
+                chemutils.atom_formal_charge_one_hot,
+                chemutils.atom_num_radical_electrons_one_hot,
+                chemutils.atom_is_aromatic_one_hot,
+                chemutils.atom_is_in_ring_one_hot,
+                chemutils.atom_chiral_tag_one_hot
+            ])
+        elif feature_mode == 'medium':
+            atom_featurizer_funs = chemutils.ConcatFeaturizer([
+                chemutils.atom_mass,
+                atom_type_one_hot, 
+                atom_bond_type_one_hot,
+                chemutils.atom_total_degree_one_hot,
+                chemutils.atom_total_num_H_one_hot,
+                chemutils.atom_is_aromatic_one_hot,
+                chemutils.atom_is_in_ring_one_hot,
+            ])
+        return chemutils.BaseAtomFeaturizer(
+        {"h": atom_featurizer_funs, 
+        "m": atom_mass_fun}
+    )
+
+    def _get_bond_featurizer(self, self_loop=True) -> dict:
+        feature_mode = self.bond_feature
+        if feature_mode == 'light':
+            return chemutils.BaseBondFeaturizer(
+                featurizer_funcs = {'e': chemutils.ConcatFeaturizer([
+                    chemutils.bond_type_one_hot
+                ])}, self_loop = self_loop
+            )
+        elif feature_mode == 'full':
+            return chemutils.CanonicalBondFeaturizer(
+                bond_data_field='e', self_loop = self_loop
+            )

mvp/data_preprocess.py

@@ -0,0 +1,78 @@
+import argparse
+from mvp.utils.preprocessing import generate_cons_spec_formulas, generate_cons_spec
+import os
+import pickle
+import pandas as pd
+from rdkit.Chem import AllChem
+from rdkit import Chem
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--spec_type", choices=('formSpec', 'binnedSpec'), required=True)
+parser.add_argument("--dataset_pth", required=True, help="path to spectra data")
+parser.add_argument("--candidates_pth", required=True, help="path to candidates data")
+parser.add_argument("--output_dir", required=True, help="path to output directory")
+parser.add_argument("--subformula_dir_pth",  default='', help="path to subformula directory if using formSpec")
+
+
+def check_args():
+
+    # create output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # check files
+    if args.spec_type == 'formSpec':
+        assert(os.path.isdir(args.subformula_dir_pth))
+    
+    assert(os.path.exists(args.dataset_pth))
+    assert(os.path.exists(args.candidates_pth))
+
+def construct_smiles_to_fp(smiles_list, r=5, fp_size=1024):
+    fpgen = AllChem.GetMorganGenerator(radius=r,fpSize=fp_size)
+    smiles_to_fp = {}
+    failed_ct = 0
+
+    for s in tqdm(smiles_list, total=len(smiles_list)):
+        try:
+            mol = Chem.MolFromSmiles(s)
+            fp = fpgen.GetFingerprint(mol)
+            smiles_to_fp[s] = fp
+        except:
+            failed_ct+=1
+    print(f'Failed to generate fingerprints for {failed_ct} smiles')
+
+    # save smiles_to_fp
+    with open(os.path.join(args.output_dir, f'morganfp_r{r}_{fp_size}.pickle'), 'wb') as f:
+        pickle.dump(smiles_to_fp, f)
+
+def construct_consensus_spectra():
+    if args.spec_type == 'formSpec':
+        df = generate_cons_spec_formulas(args.dataset_pth, args.subformula_dir_pth, args.output_dir)
+    elif args.spec_type == 'binnedSpec':
+        df = generate_cons_spec(args.dataset_pth, args.output_dir)
+
+    # save consensus spectra df
+    with open(os.path.join(args.output_dir, f'consensus_{args.spec_type}.pkl'), 'wb') as f:
+        pickle.dump(df, f)
+
+def main(data):
+
+    # generate fingerpints
+    print("Processing fingerprints...")
+    unique_smiles = data['smiles'].unique().tolist()
+    construct_smiles_to_fp(unique_smiles)
+
+    # generate consensus spectra
+    print("Processring consensus spectra...")
+    construct_consensus_spectra()
+
+
+if __name__ == '__main__':
+    args = parser.parse_args([] if "__file__" not in globals() else None)
+
+    check_args()
+
+    # load data
+    data = pd.read_csv(args.dataset_pth, sep='\t')
+
+    main(data)

mvp/definitions.py

@@ -0,0 +1,23 @@
+"""Global variables used across the package."""
+import pathlib
+
+# Dirs
+ROOT_DIR = pathlib.Path(__file__).parent.absolute()
+REPO_DIR = ROOT_DIR.parent
+DATA_DIR = REPO_DIR / 'data'
+TEST_RESULTS_DIR = REPO_DIR / 'experiments'
+ASSETS_DIR = REPO_DIR / 'assets'
+
+# C
+CHEM_ELEMS_SMALL = ['H', 'C',  'O', 'N', 'P', 'S', 'Cl', 'F', 'Br', 'I', 'B', 'As', 'Si', 'Se']
+
+MSGYM_FORMULA_VECTOR_NORM = [102.0, 59.0, 25.0, 13.0, 3.0, 6.0, 6.0, 17.0, 4.0, 4.0, 1.0, 1.0, 5.0, 2.0]
+
+#MSGYM standardization
+MSGYM_STANDARD_MH = {
+    'mz_mean': 195.155185,
+    'mz_std':127.591549
+}
+MSGYM_STANDARD_all = { # got these from Yinkai
+"mz_mean": 80.88304948022557,
+"mz_std" : 197.4588028571758}

mvp/models/__init__.py

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

mvp/models/contrastive.py

@@ -0,0 +1,416 @@
+import typing as T
+import torch
+import torch.nn as nn
+import pandas as pd
+from collections import defaultdict
+import numpy as np
+import os
+from massspecgym.models.retrieval.base import RetrievalMassSpecGymModel
+from massspecgym.models.base import Stage
+from massspecgym import utils
+
+from mvp.utils.loss import contrastive_loss, cand_spec_sim_loss, fp_loss, cons_spec_loss
+import mvp.utils.models as model_utils
+
+import torch.nn.functional as F
+
+
+class ContrastiveModel(RetrievalMassSpecGymModel):
+    def __init__(
+        self,
+        external_test = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.save_hyperparameters()
+        self.external_test = external_test
+
+        if 'use_fp' not in self.hparams:
+            self.hparams.use_fp = False
+
+        if 'loss_strategy' not in self.hparams:
+            self.hparams.loss_strategy = 'static'
+            self.hparams.contr_wt = 1.0
+            self.hparams.use_contr = True
+
+        self.spec_enc_model = model_utils.get_spec_encoder(self.hparams.spec_enc, self.hparams)
+        self.mol_enc_model = model_utils.get_mol_encoder(self.hparams.mol_enc, self.hparams)
+        
+        if self.hparams.pred_fp:
+            self.fp_loss = fp_loss(self.hparams.fp_loss_type)
+            self.fp_pred_model = model_utils.get_fp_pred_model(self.hparams)
+        if self.hparams.use_cons_spec:
+            self.cons_spec_enc_model = model_utils.get_spec_encoder(self.hparams.spec_enc, self.hparams)
+            self.cons_loss = cons_spec_loss(self.hparams.cons_loss_type)
+            
+        self.spec_view = self.hparams.spectra_view
+        
+        # result storage for testing results
+        self.result_dct = defaultdict(lambda: defaultdict(list))
+                    
+    def forward(self, batch, stage):
+        g = batch['cand'] if stage == Stage.TEST else batch['mol']
+        
+        if self.hparams.use_cons_spec and stage != Stage.TEST:
+            spec = batch['cons_spec']
+            n_peaks = batch['cons_n_peaks'] if 'cons_n_peaks' in batch else None
+            spec_enc = self.cons_spec_enc_model(spec, n_peaks)
+        else:
+            spec = batch[self.spec_view]
+            n_peaks = batch['n_peaks'] if 'n_peaks' in batch else None
+            spec_enc = self.spec_enc_model(spec, n_peaks)
+
+        fp = batch['fp'] if self.hparams.use_fp else None
+        mol_enc = self.mol_enc_model(g, fp=fp)
+
+        return spec_enc, mol_enc
+
+    def compute_loss(self, batch: dict, spec_enc, mol_enc, output):
+        loss = 0
+        losses = {}
+        contr_loss, cong_loss, noncong_loss = contrastive_loss(spec_enc, mol_enc, self.hparams.contr_temp)
+        contr_loss = self.loss_wts['contr_wt'] *contr_loss
+        losses['contr_loss'] = contr_loss.detach().item()
+        losses['cong_loss'] = cong_loss.detach().item()
+        losses['noncong_loss'] = noncong_loss.detach().item()
+        
+        loss+=contr_loss
+        if self.hparams.pred_fp:
+            fp_loss_val = self.loss_wts['fp_wt'] *self.fp_loss(output['fp'], batch['fp'])
+            loss+= fp_loss_val
+            losses['fp_loss'] = fp_loss_val.detach().item()
+
+        if 'aug_cand_enc' in output:  
+            aug_cand_loss = self.loss_wts['aug_cand_wt'] * cand_spec_sim_loss(spec_enc, output['aug_cand_enc'])
+            loss+= aug_cand_loss
+            losses['aug_cand_loss'] = aug_cand_loss.detach().item()
+        
+        if 'ind_spec' in output:
+            spec_loss = self.loss_wts['cons_spec_wt'] * self.cons_loss(spec_enc, output['ind_spec'])
+            loss+=spec_loss
+            losses['cons_spec_loss'] = spec_loss.detach().item()
+
+        losses['loss'] = loss 
+
+        return losses
+    
+    def step(
+        self, batch: dict, stage= Stage.NONE):
+        
+        # Compute spectra and mol encoding
+        spec_enc, mol_enc = self.forward(batch, stage)
+
+        if stage == Stage.TEST:
+            return dict(spec_enc=spec_enc, mol_enc=mol_enc)
+        
+        # Aux tasks
+        output = {}
+        if self.hparams.pred_fp:
+            output['fp'] = self.fp_pred_model(mol_enc)
+        
+        if self.hparams.use_cons_spec:
+            spec = batch[self.spec_view]
+            n_peaks = batch['n_peaks'] if 'n_peaks' in batch else None
+            output['ind_spec'] = self.spec_enc_model(spec, n_peaks)
+
+        # Calculate loss
+        losses = self.compute_loss(batch, spec_enc, mol_enc, output)
+
+        return losses
+
+    def on_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage) -> None:
+        # total loss
+        self.log(
+            f'{stage.to_pref()}loss',
+            outputs['loss'],
+            batch_size=len(batch['identifier']),
+            sync_dist=True,
+            prog_bar=True,
+            on_epoch=True,
+            # on_step=True
+        )
+
+        # contr loss
+        if self.hparams.use_contr:
+            self.log(
+                f'{stage.to_pref()}contr_loss',
+                outputs['contr_loss'],
+                batch_size=len(batch['identifier']),
+                sync_dist=True,
+                prog_bar=False,
+                on_epoch=True,
+                # on_step=True
+            )
+
+            # noncongruent pairs
+            self.log(
+                f'{stage.to_pref()}noncong_loss',
+                outputs['noncong_loss'],
+                batch_size=len(batch['identifier']),
+                sync_dist=True,
+                prog_bar=False,
+                on_epoch=True,
+                # on_step=True
+            )
+            
+            # congruent pairs
+            self.log(
+                f'{stage.to_pref()}cong_loss',
+                outputs['cong_loss'],
+                batch_size=len(batch['identifier']),
+                sync_dist=True,
+                prog_bar=False,
+                on_epoch=True,
+                # on_step=True
+            )
+
+
+        if self.hparams.pred_fp:
+
+            self.log(
+                f'{stage.to_pref()}_fp_loss',
+                outputs['fp_loss'],
+                batch_size=len(batch['identifier']),
+                sync_dist=True,
+                prog_bar=False,
+                on_epoch=True,
+            )
+        
+        if self.hparams.use_cons_spec:
+            self.log(
+                f'{stage.to_pref()}cons_loss',
+                outputs['cons_spec_loss'],
+                batch_size=len(batch['identifier']),
+                sync_dist=True,
+                prog_bar=False,
+                on_epoch=True,
+            )
+    
+    def test_step(self, batch, batch_idx):
+        # Unpack inputs
+        identifiers = batch['identifier']
+        cand_smiles = batch['cand_smiles']
+        id_to_ct = defaultdict(int)
+        for i in identifiers: id_to_ct[i]+=1
+        batch_ptr = torch.tensor(list(id_to_ct.values()))
+
+        outputs = self.step(batch, stage=Stage.TEST)
+        spec_enc = outputs['spec_enc']
+        mol_enc = outputs['mol_enc']
+
+        # Calculate scores
+        indexes = utils.batch_ptr_to_batch_idx(batch_ptr)
+        
+        scores = nn.functional.cosine_similarity(spec_enc, mol_enc)
+        scores = torch.split(scores, list(id_to_ct.values()))
+
+        cand_smiles = utils.unbatch_list(batch['cand_smiles'], indexes)
+        labels = utils.unbatch_list(batch['label'], indexes)
+        
+        return dict(identifiers=list(id_to_ct.keys()), scores=scores, cand_smiles=cand_smiles, labels=labels)
+    
+    def on_test_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage = Stage.TEST) -> None:
+        
+        # save scores
+        for i, cands, scores, l in zip(outputs['identifiers'], outputs['cand_smiles'], outputs['scores'], outputs['labels']):
+            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])
+            
+    def _compute_rank(self, scores, labels):
+        if not any(labels):
+            return -1
+        scores = np.array(scores)
+        target_score = scores[labels][0]
+        rank = np.count_nonzero(scores >=target_score)
+        return rank
+    
+    def on_test_epoch_end(self) -> None:
+
+        self.df_test = pd.DataFrame.from_dict(self.result_dct, orient='index').reset_index().rename(columns={'index': 'identifier'})
+
+        # Compute rank
+        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]:
+        monitors = [
+            {"monitor": f"{Stage.TRAIN.to_pref()}loss", "mode": "min", "early_stopping": False}, # monitor train loss
+        ]
+        return monitors
+    
+    def _update_loss_weights(self)-> None:
+        if self.hparams.loss_strategy == 'linear':
+            for loss in self.loss_wts:
+                self.loss_wts[loss] += self.loss_updates[loss]
+        elif self.hparams.loss_strategy == 'manual':
+            for loss in self.loss_wts:
+                if self.current_epoch in self.loss_updates[loss]:
+                    self.loss_wts[loss] = self.loss_updates[loss][self.current_epoch]
+
+    def on_train_epoch_end(self) -> None:
+        self._update_loss_weights()      
+    
+class MultiViewContrastive(ContrastiveModel):
+
+    def __init__(self,
+                 **kwargs):
+        
+        super().__init__(**kwargs)
+    
+        # build fingerprint encoder model
+        if self.hparams.use_fp: 
+            self.fp_enc_model = model_utils.get_fp_enc_model(self.hparams)
+        
+        # build NL encoder model
+        # if self.hparams.use_NL_spec:
+        #     self.NL_enc_model = model_utils.get_spec_encoder(self.hparams.spec_enc, self.hparams)
+            
+    def forward(self, batch, stage):
+        g = batch['cand'] if stage == Stage.TEST else batch['mol']
+        
+        spec = batch[self.spec_view]
+        n_peaks = batch['n_peaks'] if 'n_peaks' in batch else None
+
+        spec_enc = self.spec_enc_model(spec, n_peaks)
+        mol_enc = self.mol_enc_model(g)
+        views = {'spec_enc': spec_enc, 'mol_enc': mol_enc}
+        
+        if self.hparams.use_fp:
+            fp_enc = self.fp_enc_model(batch['fp'])
+            views['fp_enc'] = fp_enc
+
+        if self.hparams.use_cons_spec:
+            spec = batch['cons_spec']
+            n_peaks = batch['cons_n_peaks'] if 'cons_n_peaks' in batch else None
+            spec_enc = self.cons_spec_enc_model(spec, n_peaks)
+            views['cons_spec_enc'] = spec_enc
+
+        if self.hparams.use_NL_spec:
+            spec = batch['NL_spec']
+            n_peaks = batch['NL_n_peaks'] if 'NL_n_peaks' in batch else None
+            spec_enc = self.NL_enc_model(spec, n_peaks) 
+            views['NL_spec_enc'] = spec_enc
+        return views
+    
+    def step(
+        self, batch: dict, stage= Stage.NONE):
+        
+        # Compute spectra and mol encoding
+        views = self.forward(batch, stage)
+
+        if stage == Stage.TEST:
+            return views
+        
+        # Calculate loss
+        losses = self.compute_loss(batch, views)
+
+        return losses
+    
+    def compute_loss(self, batch: dict, views: dict):
+        loss = 0
+        losses = {}
+        for v1, v2 in self.hparams.contr_views:
+            contr_loss, cong_loss, noncong_loss = contrastive_loss(views[v1], views[v2], self.hparams.contr_temp)
+            loss+=contr_loss
+
+            losses[f'{v1[:-4]}-{v2[:-4]}_contr_loss'] = contr_loss.detach().item()
+            losses[f'{v1[:-4]}-{v2[:-4]}_cong_loss'] = cong_loss.detach().item()
+            losses[f'{v1[:-4]}-{v2[:-4]}_noncong_loss'] = noncong_loss.detach().item()
+        
+        losses['loss'] = loss 
+
+        return losses
+    
+    def on_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage) -> None:
+        # total loss
+        self.log(
+            f'{stage.to_pref()}loss',
+            outputs['loss'],
+            batch_size=len(batch['identifier']),
+            sync_dist=True,
+            prog_bar=True,
+            on_epoch=True,
+            # on_step=True
+        )
+
+        for v1, v2 in self.hparams.contr_views:
+            self.log(
+            f'{stage.to_pref()}{v1[:-4]}-{v2[:-4]}_contr_loss',
+            outputs[f'{v1[:-4]}-{v2[:-4]}_contr_loss'],
+            batch_size=len(batch['identifier']),
+            sync_dist=True,
+            on_epoch=True,
+        )
+            self.log(
+            f'{stage.to_pref()}{v1[:-4]}-{v2[:-4]}_cong_loss',
+            outputs[f'{v1[:-4]}-{v2[:-4]}_cong_loss'],
+            batch_size=len(batch['identifier']),
+            sync_dist=True,
+            on_epoch=True,
+        )
+            self.log(
+            f'{stage.to_pref()}{v1[:-4]}-{v2[:-4]}_noncong_loss',
+            outputs[f'{v1[:-4]}-{v2[:-4]}_noncong_loss'],
+            batch_size=len(batch['identifier']),
+            sync_dist=True,
+            on_epoch=True,
+        )
+            
+    def test_step(self, batch, batch_idx):
+        # Unpack inputs
+        identifiers = batch['identifier']
+        cand_smiles = batch['cand_smiles']
+        id_to_ct = defaultdict(int)
+        for i in identifiers: id_to_ct[i]+=1
+        batch_ptr = torch.tensor(list(id_to_ct.values()))
+
+        outputs = self.step(batch, stage=Stage.TEST)
+        scores = {}
+        for v1, v2 in self.hparams.contr_views:
+            # if 'cons_spec_enc' in (v1, v2):
+            #     continue
+            v1_enc = outputs[v1]
+            v2_enc = outputs[v2]
+            
+            s = nn.functional.cosine_similarity(v1_enc, v2_enc)
+            scores[f'{v1[:-4]}-{v2[:-4]}_scores'] = torch.split(s, list(id_to_ct.values()))
+
+        indexes = utils.batch_ptr_to_batch_idx(batch_ptr)
+
+        cand_smiles = utils.unbatch_list(batch['cand_smiles'], indexes)
+        labels = utils.unbatch_list(batch['label'], indexes)
+        
+        return dict(identifiers=list(id_to_ct.keys()), scores=scores, cand_smiles=cand_smiles, labels=labels)
+    
+    def on_test_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage = Stage.TEST) -> None:
+
+        # save scores
+        for i, cands, l in zip(outputs['identifiers'], outputs['cand_smiles'], outputs['labels']):
+            self.result_dct[i]['candidates'].extend(cands)
+            self.result_dct[i]['labels'].extend([x.cpu().item() for x in l])
+
+        for v1, v2 in self.hparams.contr_views:
+            for i, scores in zip(outputs['identifiers'], outputs['scores'][f'{v1[:-4]}-{v2[:-4]}_scores']):
+                self.result_dct[i][f'{v1[:-4]}-{v2[:-4]}_scores'].extend(scores.cpu().tolist())
+
+    def _get_top_cand(self, scores, candidates):
+        return candidates[np.argmax(np.array(scores))]
+
+    def on_test_epoch_end(self) -> None:
+
+        self.df_test = pd.DataFrame.from_dict(self.result_dct, orient='index').reset_index().rename(columns={'index': 'identifier'})
+
+        # Compute rank
+        if not self.external_test:
+            for v1, v2 in self.hparams.contr_views:
+                self.df_test[f'{v1[:-4]}-{v2[:-4]}_rank'] = self.df_test.apply(lambda row: self._compute_rank(row[f'{v1[:-4]}-{v2[:-4]}_scores'], row['labels']), axis=1)
+        
+        if self.external_test:
+            self.df_test.drop('labels', axis=1, inplace=True)
+            for v1, v2 in self.hparams.contr_views:
+                self.df_test[f'top_{v1[:-4]}-{v2[:-4]}_cand'] = self.df_test.apply(lambda row: self._get_top_cand(row[f'{v1[:-4]}-{v2[:-4]}_scores'], row['candidates']), axis=1)
+        self.df_test.to_pickle(self.df_test_path)

mvp/models/encoders.py

@@ -0,0 +1,29 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MLP(nn.Module):
+    def __init__(self, in_dim, hidden_dims, dropout=0.1, final_activation=None):
+        super(MLP, self).__init__()
+
+        self.dropout = nn.Dropout(dropout)
+        self.has_final_activation = False
+        layers = [nn.Linear(in_dim, hidden_dims[0])]
+        for d1, d2 in zip(hidden_dims[:-1], hidden_dims[1:]):
+            layers.append(nn.Linear(d1, d2))
+        self.layers = nn.ModuleList(layers)
+        if final_activation is not None:
+            self.has_final_activation = True
+    
+            self.final_activation = {'relu': F.relu,
+                                    'sigmoid': F.sigmoid,
+                                    'softmax': F.softmax,}[final_activation]
+    
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+            if i < len(self.layers) -1:
+                x = F.relu(x)
+                x = self.dropout(x)
+            elif self.has_final_activation:
+                x = self.final_activation(x)
+        return x

mvp/models/mol_encoder.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+import dgl
+from dgllife.model import GCN, GAT
+
+class MolEnc(nn.Module):
+
+    def __init__(self,
+                 args,
+                 in_dim,):
+        super().__init__()
+
+        self.return_emb = False
+
+        if args.model in ('crossAttenContrastive', 'filipContrastive'):
+            self.return_emb = True
+
+        dropout = [args.gnn_dropout for _ in range(len(args.gnn_channels))]
+        batchnorm = [True for _ in range(len(args.gnn_channels))]
+        gnn_map = {
+            "gcn": GCN(in_dim, args.gnn_channels, batchnorm = batchnorm, dropout = dropout),
+            "gat": GAT(in_dim, args.gnn_channels, args.attn_heads)
+        }
+        self.GNN = gnn_map[args.gnn_type]
+        self.pool = dgl.nn.pytorch.glob.MaxPooling()
+        
+        if not self.return_emb:
+            self.fc1_graph = nn.Linear(args.gnn_channels[len(args.gnn_channels) - 1], args.gnn_hidden_dim * 2)
+            self.fc2_graph = nn.Linear(args.gnn_hidden_dim * 2, args.final_embedding_dim)
+
+            self.dropout = nn.Dropout(args.fc_dropout)
+            self.relu = nn.ReLU()
+
+    def forward(self, g, fp=None) -> torch.Tensor:
+        g1 = g
+        f1 = g.ndata['h']
+
+        f = self.GNN(g1, f1)
+        if self.return_emb:
+            return f
+        h = self.pool(g1, f)
+        if fp is not None:
+            h = torch.concat((h, fp), dim=-1)
+        h1 = self.relu(self.fc1_graph(h))
+        h1 = self.dropout(h1)
+        h1 = self.fc2_graph(h1)
+        h1 = self.dropout(h1)
+
+        return h1
+    

mvp/models/spec_encoder.py

@@ -0,0 +1,85 @@
+import torch.nn as nn
+import torch
+from mvp.models.encoders import MLP
+from torch_geometric.nn import global_mean_pool
+
+
+class SpecEncMLP_BIN(nn.Module):
+    def __init__(self, args, out_dim=None):
+        super(SpecEncMLP_BIN, self).__init__()
+
+        if not out_dim:
+            out_dim = args.final_embedding_dim
+
+        bin_size = int(args.max_mz / args.bin_width)
+        self.dropout = nn.Dropout(args.fc_dropout)
+        self.mz_fc1 = nn.Linear(bin_size, out_dim * 2)
+        self.mz_fc2 = nn.Linear(out_dim* 2, out_dim * 2)
+        self.mz_fc3 = nn.Linear(out_dim * 2, out_dim)
+        self.relu = nn.ReLU()
+    
+    def forward(self, mzi_b, n_peaks=None):
+                
+       h1 = self.mz_fc1(mzi_b)
+       h1 = self.relu(h1)
+       h1 = self.dropout(h1)
+       h1 = self.mz_fc2(h1)
+       h1 = self.relu(h1)
+       h1 = self.dropout(h1)
+       mz_vec = self.mz_fc3(h1)
+       mz_vec = self.dropout(mz_vec)
+       
+       return mz_vec
+
+    
+class SpecFormulaTransformer(nn.Module):
+    def __init__(self, args, out_dim=None):
+        super(SpecFormulaTransformer, self).__init__()
+        in_dim = len(args.element_list)
+        if args.add_intensities: # intensity
+            in_dim+=1
+        if args.spectra_view == "SpecFormulaMz": #mz
+            in_dim+=1 
+
+        self.returnEmb = False
+        
+        self.formulaEnc = MLP(in_dim=in_dim, hidden_dims=args.formula_dims, dropout=args.formula_dropout)
+        
+        self.use_cls = args.use_cls
+        if args.use_cls:
+            self.cls_embed = torch.nn.Embedding(1,args.formula_dims[-1])
+        encoder_layer = nn.TransformerEncoderLayer(d_model=args.formula_dims[-1], nhead=2, batch_first=True)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
+
+        if not out_dim:
+            out_dim = args.final_embedding_dim
+        self.fc = nn.Linear(args.formula_dims[-1], out_dim)
+        
+    def forward(self, spec, n_peaks):
+        h = self.formulaEnc(spec)
+        pad = (spec == -5)
+        pad = torch.all(pad, -1)
+
+        if self.use_cls:
+            cls_embed = self.cls_embed(torch.tensor(0).to(spec.device))
+            h = torch.concat((cls_embed.repeat(spec.shape[0], 1).unsqueeze(1), h), dim=1)
+            pad = torch.concat((torch.tensor(False).repeat(pad.shape[0],1).to(spec.device), pad), dim=1)
+            h = self.transformer(h, src_key_padding_mask=pad)
+            h = h[:,0,:]
+        else:
+            h = self.transformer(h, src_key_padding_mask=pad)
+
+            if self.returnEmb:
+                # repad h
+                h[pad] = -5
+                return h
+            
+            h = h[~pad].reshape(-1, h.shape[-1])
+            indecies = torch.tensor([i for i, count in enumerate(n_peaks) for _ in range(count)]).to(h.device)
+            h = global_mean_pool(h, indecies)
+            
+        h = self.fc(h)
+
+        return h
+    
+

mvp/params_binnedSpec.yaml

@@ -0,0 +1,122 @@
+
+# Experiment setup 
+job_key: ''
+run_name: 'binnedSpec_experiment'
+run_details: ""
+project_name: ''
+wandb_entity_name: 'mass-spec-ml'
+no_wandb: True
+seed: 0
+debug: False
+checkpoint_pth: ""
+
+# Training setup
+max_epochs: 1000 
+accelerator: 'gpu'
+devices: [1]
+log_every_n_steps: 250
+val_check_interval: 1.0
+
+# Data paths
+candidates_pth: ../data/sample/candidates_mass.json
+dataset_pth: "../data/sample/data.tsv" 
+subformula_dir_pth: ""
+split_pth: 
+fp_dir_pth: '../data/sample/morganfp_r5_1024.pickle'
+cons_spec_dir_pth: "../data/sample/consensus_binnedSpec.pkl"
+NL_spec_dir_pth: ""
+partial_checkpoint: ""
+
+# General hyperparameters
+batch_size: 64
+lr: 5.0e-4
+weight_decay: 0
+contr_temp: 0.05
+early_stopping_patience: 300
+loss_strategy: 'static' # static, linear, manual
+num_workers: 50
+
+
+############################## Data transforms ##############################
+# - Spectra 
+spectra_view:  SpecBinnerLog 
+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: True
+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: MLP_BIN 
+mol_enc: "GNN"
+model:  "MultiviewContrastive" 
+contr_views: [['spec_enc', 'mol_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: True
+
+# -- 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/params_formSpec.yaml

@@ -0,0 +1,121 @@
+# Experiment setup 
+job_key: ''
+run_name: 'filip_large'
+run_details: ""
+project_name: ''
+wandb_entity_name: 'mass-spec-ml'
+no_wandb: True
+seed: 0
+debug: False
+checkpoint_pth: #'../pretrained_models/msgym_formSpec.ckpt'
+
+# Training setup
+max_epochs: 2000
+accelerator: 'gpu'
+devices: [1]
+log_every_n_steps: 250
+val_check_interval: 1.0
+
+# Data paths
+candidates_pth: ../data/sample/candidates_mass.json
+dataset_pth: ../data/MassSpecGym/data/sample_data.tsv
+subformula_dir_pth: ../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"
+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
+#  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
+mol_enc: "GNN"
+model:  MultiviewContrastive
+contr_views: [['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

@@ -0,0 +1,9 @@
+# 1. preprocess data (subformula labels should be obtained through MIST)
+python subformula_assign/assign_subformulae.py --spec-files ../data/sample/data.tsv --output-dir ../data/sample/subformulae_default --max-formulae 60 --labels-file ../data/sample/data.tsv
+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

mvp/test.py

@@ -0,0 +1,117 @@
+import argparse
+import datetime
+import sys
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from rdkit import RDLogger
+import pytorch_lightning as pl
+from pytorch_lightning import Trainer
+from massspecgym.models.base import Stage
+import os
+
+from mvp.data.data_module import TestDataModule
+from mvp.data.datasets import ContrastiveDataset
+from mvp.utils.data import get_spec_featurizer, get_mol_featurizer, get_test_ms_dataset
+from mvp.utils.models import get_model
+
+from mvp.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)
+def main(params):
+    # Seed everything
+    pl.seed_everything(params['seed'])
+        
+    # Init paths to data files
+    if params['debug']:
+        params['dataset_pth'] = "../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)
+
+    # Init data module
+    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']
+    )
+
+    model = get_model(params['model'], params)
+    model.df_test_path = params['df_test_path']
+    
+    # Init trainer
+    trainer = Trainer(
+        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
+    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
+    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):
+            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
+                break
+    assert(params['checkpoint_pth'] != '')
+    
+    if args.candidates_pth:
+        params['candidates_pth'] = args.candidates_pth
+    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")
+        
+    main(params)

mvp/train.py

@@ -0,0 +1,137 @@
+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
+import pytorch_lightning as pl
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+
+
+from mvp.data.data_module import ContrastiveDataModule
+
+from mvp.definitions import TEST_RESULTS_DIR
+import yaml
+from mvp.data.datasets import ContrastiveDataset
+from functools import partial
+
+from mvp.utils.data import get_ms_dataset, get_spec_featurizer, get_mol_featurizer
+from mvp.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")
+
+def main(params):
+    # Seed everything
+    pl.seed_everything(params['seed'])
+        
+    # Init paths to data files
+    if params['debug']:
+        params['dataset_pth'] = "../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'], mask_peak_ratio=params['mask_peak_ratio'], aug_cands=params['aug_cands'])
+    data_module = ContrastiveDataModule(
+        dataset=dataset,
+        collate_fn=collate_fn,
+        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
+        )
+
+    # Init callbacks for checkpointing and early stopping
+    callbacks = [pl.callbacks.ModelCheckpoint(save_last=False) ]
+    for i, monitor in enumerate(model.get_checkpoint_monitors()):
+        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}',
+            auto_insert_metric_name=True,
+            save_last=(i == 0)
+        )
+        callbacks.append(checkpoint)
+        if monitor.get('early_stopping', False):
+            early_stopping = EarlyStopping(
+                monitor=monitor_name,
+                mode=monitor['mode'],
+                verbose=True,
+                patience=params['early_stopping_patience'],
+            )
+            callbacks.append(early_stopping)
+
+    # Init trainer
+    trainer = Trainer(
+        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'],
+        callbacks=callbacks,
+        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
+    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
+
+    if not params['df_test_path']:
+        params['df_test_path'] = os.path.join(experiment_dir, "result.pkl")
+
+    main(params)

mvp/utils/__init__.py

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

mvp/utils/data.py

@@ -0,0 +1,212 @@
+import os
+import json
+import numpy as np
+
+from mvp.data.transforms import SpecBinner, SpecBinnerLog, SpecFormulaFeaturizer
+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
+import matchms
+
+class Subformula_Loader:
+    def __init__(self, spectra_view, dir_path) -> None:
+
+        self.dir_path = dir_path
+        if spectra_view == 'SpecFormula':
+            self.load = self.load_subformula_data
+        elif spectra_view == "SpecFormulaMz":
+            self.load = self.load_subformula_dict
+        else:
+            raise Exception("Spectra view is not supported.")
+
+    def __call__(self, ids):
+        id_to_form_spec = {}
+        for id in ids:
+            data = self.load(id)
+            if data:
+                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 
+        '''
+        try:
+            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'])
+
+            # 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_dict(self, spec_id: str):
+        '''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)
+            mzs = np.array(data['output_tbl']['mz'])
+            formulas = np.array(data['output_tbl']['formula'])
+            intensities = np.array(data['output_tbl']['ms2_inten'])
+
+            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]
+            intensities = intensities[ind]
+            return {'formulas': mz_to_formulas, 'formula_mzs': mzs, 'formula_intensities': intensities}
+        except:
+            return None
+
+def make_tmp_subformula_spectra(row):
+        return {'formulas':[row['formula']], 'formula_mzs':[float(row['precursor_mz'])], 'formula_intensities':[1.0]}
+
+def get_spec_featurizer(spectra_view: T.Union[str, list[str]],
+                         params) -> T.Union[SpecTransform, T.Dict[str, SpecTransform]]:
+    
+    featurizers = {"BinnedSpectra": SpecBinner,
+        "SpecBinnerLog": SpecBinnerLog,
+        "SpecFormula": SpecFormulaFeaturizer}
+
+    spectra_featurizer = {}
+
+    if isinstance(spectra_view, str):
+        spectra_view = [spectra_view]
+
+    for view in spectra_view:
+        featurizer_params = {'max_mz': params['max_mz']}
+        if view in ["BinnedSpectra", "SpecBinnerLog"]:
+            featurizer_params.update({'bin_width': params['bin_width']})
+        elif view in ["SpecFormula"]:
+            featurizer_params.update({'element_list': params['element_list'], 'add_intensities': params['add_intensities'], 'formula_normalize_vector': MSGYM_FORMULA_VECTOR_NORM})
+        
+        spectra_featurizer[view] = featurizers[view](**featurizer_params)
+
+    return spectra_featurizer
+
+def get_mol_featurizer(molecule_view: T.Union[str, T.List[str]], params) -> MolTransform:
+    featurizes = {'MolGraph':MolToGraph}
+    mol_featurizer = {}
+
+    if isinstance(molecule_view, str):
+        molecule_view = [molecule_view]
+    for view in molecule_view:
+        featurizer_params = {}
+        if view in ('MolGraph'):
+            featurizer_params.update({'atom_feature': params['atom_feature'], 'bond_feature': params['bond_feature'], 'element_list': params['element_list']})
+        
+        if len(molecule_view) == 1:
+            return featurizes[view](**featurizer_params)
+
+        mol_featurizer[view] = featurizes[view](**featurizer_params)
+    
+    return mol_featurizer
+
+def get_test_ms_dataset(spectra_view: T.Union[str, T.List[str]],
+                 mol_view: T.Union[str, T.List[str]],
+                 spectra_featurizer: SpecTransform,
+                 mol_featurizer: MolTransform,
+                 params,
+                external_test: bool = False,):
+    
+    use_formulas = False
+    
+    views = []
+    for v in [spectra_view, mol_view]:
+        if isinstance(v, str):
+            views.append(v)
+        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']}
+    if "SpecFormula" in views or "SpecFormulaMz" in views:
+        dataset_params.update({'subformula_dir_pth': params['subformula_dir_pth']})
+        use_formulas = True 
+    
+    if params['use_cons_spec']:
+        dataset_params.update({'cons_spec_dir_pth': params['cons_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']})
+
+    return jestr_datasets.ExpandedRetrievalDataset(use_formulas=use_formulas, external_test=external_test, **dataset_params)
+    
+def get_ms_dataset(spectra_view: str,
+                 mol_view: str,
+                 spectra_featurizer: SpecTransform,
+                 mol_featurizer: MolTransform,
+                 params):
+    
+
+    # set up dataset_parameters
+    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']})
+        use_formulas = True
+
+    if params['pred_fp'] or params['use_fp']:
+        dataset_params.update({'fp_dir_pth': params['fp_dir_pth']})
+        
+    if params['use_cons_spec']:
+        dataset_params.update({'cons_spec_dir_pth': params['cons_spec_dir_pth']})
+
+    # select dataset
+    if params['aug_cands']:
+        return jestr_datasets.MassSpecDataset_Candidates(**dataset_params)
+    elif use_formulas:
+        return jestr_datasets.MassSpecDataset_PeakFormulas(**dataset_params)
+    
+    return jestr_datasets.JESTR1_MassSpecDataset(**dataset_params)
+
+class PrepMatchMS:
+    def __init__(self, spectra_view) -> None:
+
+        if spectra_view == 'SpecFormula':
+            self.prepare = self.specFormula
+        elif spectra_view == "SpecFormulaMz":
+            self.prepare = self.specFormulaMz
+        elif spectra_view in ('SpecBinnerLog', 'BinnedSpectra', 'SpecMzIntTokenizer'):
+            self.prepare = self.specMzInt
+        else:
+            raise Exception("Spectra view is not supported.")
+        
+    def specFormulaMz(self, row):
+        
+        return matchms.Spectrum(
+            mz = np.array([float(m) for m in row["mzs"].split(",")]),
+            intensities = np.array(
+                        [float(i) for i in row["intensities"].split(",")]
+                    ),
+            metadata = {'precursor_mz': row['precursor_mz'], 'formulas': row['formulas']}
+        )
+    
+    def specFormula(self, row):
+
+        return matchms.Spectrum(
+            mz = np.array(row['formula_mzs']),
+            intensities = np.array(row['formula_intensities']),
+            metadata = {'precursor_mz': row['precursor_mz'], 'formulas': np.array(row['formulas']), 'precursor_formula': row['precursor_formula']}
+        )
+    
+    def specMzInt(self, row):
+        return matchms.Spectrum(
+            mz = row['mzs'],
+            intensities = row['intensities'],
+            metadata = {'precursor_mz': row['precursor_mz']}
+        )

mvp/utils/eval.py

@@ -0,0 +1,157 @@
+from MassSpecGym.massspecgym.utils import MyopicMCES
+import numpy as np
+import tqdm
+from multiprocessing import Pool
+
+import os
+import pandas as pd
+
+class Compute_Myopic_MCES:
+    mces_compute = MyopicMCES()
+    
+
+    def compute_mces(tar_cand):
+        target, cand = tar_cand
+        
+        dist = Compute_Myopic_MCES.mces_compute(target, cand)
+        return (tar_cand, dist)
+    
+    def compute_mces_parallel(target_cand_list, n_processes=25):
+
+
+        with Pool(processes=n_processes) as pool:
+            results = list(tqdm.tqdm(pool.imap(Compute_Myopic_MCES.compute_mces, target_cand_list), total=len(target_cand_list)))
+        return results
+
+class Compute_Myopic_MCES_timeout:
+    mces_compute = MyopicMCES()
+
+    @staticmethod
+    def compute_mces(tar_cand):
+        target, cand = tar_cand
+        dist = Compute_Myopic_MCES.mces_compute(target, cand)
+        return (tar_cand, dist)
+
+    @staticmethod
+    def compute_mces_parallel(target_cand_list, n_processes=35, timeout=60):  # timeout in seconds
+        results = []
+
+        with Pool(processes=n_processes) as pool:
+            async_results = [
+                pool.apply_async(Compute_Myopic_MCES.compute_mces, args=(tar_cand,))
+                for tar_cand in target_cand_list
+            ]
+            for async_res in tqdm.tqdm(async_results, total=len(target_cand_list)):
+                try:
+                    result = async_res.get(timeout=timeout)
+                except Exception as e:
+                    # You can log the error or return a default value
+                    result = (None, f"Timeout or error")
+                results.append(result)
+
+        return results
+
+    
+def get_result_files(exp_dir, spec_type, views_type):
+    files = os.listdir(exp_dir)
+    mass_result = ''
+    form_result = ''
+
+    for f in files:
+        try:
+            _, s, views = f.split('_')
+        except:
+            continue
+        
+        if s == spec_type and views == views_type:
+            print(exp_dir / f)
+
+            files = os.listdir(exp_dir / f)
+            for fr in files:
+                if 'mass_result' in fr:
+                    mass_result = exp_dir / f / fr
+                elif 'result' in fr:
+                    form_result = exp_dir / f/ fr
+            
+    return mass_result, form_result
+
+# get target
+def get_target(candidates, labels):
+    return np.array(candidates)[labels][0]
+
+# get mol rank at 1
+def get_top_cand(candidates, scores):
+    return candidates[np.argmax(scores)]
+
+# split into hit rates
+def convert_rank_to_hit_rates(row, rank_col ,top_k=[1,5,20]):
+    top_k_hits ={}
+    rank = row[rank_col]
+    for k in top_k:
+        if rank <= k:
+            top_k_hits[f'{rank_col}-hit_rate@{k}'] = 1
+        else:
+            top_k_hits[f'{rank_col}-hit_rate@{k}'] = 0
+    return pd.Series(top_k_hits)
+
+#################### Rank aggregation #######################
+from collections import defaultdict
+import numpy as np
+from scipy.stats import rankdata
+
+def borda_count(candidates, score_lists, target):
+    scores = defaultdict(int)
+    N = len(candidates)
+    for score_list in score_lists:
+        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
+        for rank, (mol, _) in enumerate(ranked_list, start=1):
+            scores[mol] += N - rank + 1
+    ranked_candidates = [mol for mol, _ in sorted(scores.items(), key=lambda x: x[1], reverse=True)]
+    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
+
+def average_rank(candidates, score_lists, target):
+    rank_sums = defaultdict(list)
+    for score_list in score_lists:
+        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
+        for rank, (mol, _) in enumerate(ranked_list, start=1):
+            rank_sums[mol].append(rank)
+    avg_ranks = {mol: np.mean(ranks) for mol, ranks in rank_sums.items()}
+    ranked_candidates = [mol for mol, _ in sorted(avg_ranks.items(), key=lambda x: x[1])]
+    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
+
+def reciprocal_rank_aggregation(candidates, score_lists, target):
+    scores = defaultdict(float)
+    for score_list in score_lists:
+        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
+        for rank, (mol, _) in enumerate(ranked_list, start=1):
+            scores[mol] += 1 / rank
+    ranked_candidates = [mol for mol, _ in sorted(scores.items(), key=lambda x: x[1], reverse=True)]
+    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
+
+def weighted_voting(candidates, score_lists, weights, target):
+    scores = defaultdict(float)
+    for weight, score_list in zip(weights, score_lists):
+        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
+        for rank, (mol, _) in enumerate(ranked_list, start=1):
+            scores[mol] += weight / rank
+    ranked_candidates = [mol for mol, _ in sorted(scores.items(), key=lambda x: x[1], reverse=True)]
+    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
+
+def median_rank(candidates, score_lists, target):
+    rank_sums = defaultdict(list)
+    for score_list in score_lists:
+        ranked_list = sorted(zip(candidates, score_list), key=lambda x: x[1], reverse=True)
+        for rank, (mol, _) in enumerate(ranked_list, start=1):
+            rank_sums[mol].append(rank)
+    median_ranks = {mol: np.median(ranks) for mol, ranks in rank_sums.items()}
+    ranked_candidates = [mol for mol, _ in sorted(median_ranks.items(), key=lambda x: x[1])]
+    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None
+
+def score_based_aggregation(candidates, score_lists, target):
+    scores = defaultdict(list)
+    for score_list in score_lists:
+        for mol, score in zip(candidates, score_list):
+            scores[mol].append(score)
+    avg_scores = {mol: np.mean(vals) for mol, vals in scores.items()}
+    ranked_candidates = [mol for mol, _ in sorted(avg_scores.items(), key=lambda x: x[1], reverse=True)]
+    return ranked_candidates.index(target) + 1 if target in ranked_candidates else None

mvp/utils/general.py

@@ -0,0 +1,29 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+def pad_graph_nodes(mol_enc, g_n_nodes):
+    """
+    Args:
+        mol_enc: 2D tensor of shape (sum_nodes, D)
+                 Node embeddings for each molecule.
+        g_n_nodes: list[int]  Number of nodes per graph (len = B)
+
+    Returns:
+        padded: (B, max_nodes, D) tensor
+        mask:   (B, max_nodes) bool tensor, True for valid nodes
+    """
+
+    # Already concatenated: shape (sum_nodes, D)
+    B = len(g_n_nodes)
+    D = mol_enc.shape[1]
+    max_nodes = max(g_n_nodes)
+    padded = mol_enc.new_zeros((B, max_nodes, D))
+    mask = torch.zeros((B, max_nodes), dtype=torch.bool, device=mol_enc.device)
+
+    idx = 0
+    for i, n in enumerate(g_n_nodes):
+        padded[i, :n] = mol_enc[idx:idx+n]
+        mask[i, :n] = True
+        idx += n
+    return padded, mask

mvp/utils/loss.py

@@ -0,0 +1,78 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def contrastive_loss(v1, v2, tau=1.0) -> torch.Tensor:
+        v1_norm = torch.norm(v1, dim=1, keepdim=True)
+        v2_norm = torch.norm(v2, dim=1, keepdim=True)
+        
+        v2T = torch.transpose(v2, 0, 1)
+        
+        inner_prod = torch.matmul(v1, v2T)
+        
+        v2_normT = torch.transpose(v2_norm, 0, 1)
+        
+        norm_mat = torch.matmul(v1_norm, v2_normT)
+        
+        loss_mat = torch.div(inner_prod, norm_mat)
+        
+        loss_mat = loss_mat * (1/tau)
+        
+        loss_mat = torch.exp(loss_mat)
+        
+        numerator = torch.diagonal(loss_mat)
+        numerator = torch.unsqueeze(numerator, 0)
+        
+        Lv1_v2_denom = torch.sum(loss_mat, dim=1, keepdim=True)
+        Lv1_v2_denom = torch.transpose(Lv1_v2_denom, 0, 1)
+        #Lv1_v2_denom = Lv1_v2_denom - numerator
+        
+        Lv2_v1_denom = torch.sum(loss_mat, dim=0, keepdim=True)
+        #Lv2_v1_denom = Lv2_v1_denom - numerator
+        
+        Lv1_v2 = torch.div(numerator, Lv1_v2_denom)
+        
+        Lv1_v2 = -1 * torch.log(Lv1_v2)
+        Lv1_v2 = torch.mean(Lv1_v2)
+        
+        Lv2_v1 = torch.div(numerator, Lv2_v1_denom)
+        
+        Lv2_v1 = -1 * torch.log(Lv2_v1)
+        Lv2_v1 = torch.mean(Lv2_v1)
+        
+        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)
+
+

mvp/utils/models.py

@@ -0,0 +1,38 @@
+from mvp.models.spec_encoder import SpecEncMLP_BIN, SpecFormulaTransformer
+from mvp.models.mol_encoder import MolEnc
+from mvp.models.encoders import MLP
+from mvp.models.contrastive import ContrastiveModel, MultiViewContrastive
+
+def get_spec_encoder(spec_enc:str, args):
+    return {"MLP_BIN": SpecEncMLP_BIN,
+            "Transformer_Formula": SpecFormulaTransformer}[spec_enc](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):
+    
+    if model == 'contrastive':
+        model= ContrastiveModel(**params)
+    elif model == "MultiviewContrastive":
+        model = MultiViewContrastive(**params)
+    else:
+        raise Exception(f"Model {model} not implemented.")
+    
+    # If checkpoint path is provided, load the model from the checkpoint instead
+    if params['checkpoint_pth'] is not None and params['checkpoint_pth'] != "":
+        model = type(model).load_from_checkpoint(
+            params['checkpoint_pth'],
+            log_only_loss_at_stages=params['log_only_loss_at_stages'],
+            df_test_path=params['df_test_path']
+        )
+        print("Loaded Model from checkpoint")
+
+    return model

mvp/utils/preprocessing.py

@@ -0,0 +1,149 @@
+import pandas as pd
+import pickle
+import numpy as np
+import mvp.utils.data as data_utils
+import collections
+import os
+import requests
+import tqdm
+from multiprocessing import Pool
+from urllib.parse import quote
+from tqdm import tqdm
+
+class NPClassProcess:
+    def process_smiles(smiles):
+        try:
+            encoded_smiles = quote(smiles)
+            url = f"https://npclassifier.gnps2.org/classify?smiles={encoded_smiles}"
+            r = requests.get(url)
+            return (smiles, r.json())
+        except:
+            return (smiles, None)
+    
+    def NPclass_from_smiles(pth, output_dir, n_processes=20):
+
+        data = pd.read_csv(pth, sep='\t')
+        unique_smiles = data['smiles'].unique().tolist()
+
+        items = unique_smiles
+
+        with Pool(processes=n_processes) as pool:
+            results = list(tqdm(pool.imap(NPClassProcess.process_smiles, items), total=len(items)))
+
+        failed_ct = 0
+        smiles_to_class = {}
+        for s, out in results:
+            if out is None:
+                smiles_to_class[s] = 'NA'
+                failed_ct+=1
+            else:
+                smiles_to_class[s] = out
+        file_pth = os.path.join(output_dir, 'SMILES_TO_CLASS.pkl')
+        with open(file_pth, 'wb') as f:
+            pickle.dump(smiles_to_class, f)
+        print(f'Failed to process {failed_ct} SMILES')
+        print(f'result file saved to {file_pth}')
+        return file_pth
+
+
+
+def construct_NL_spec(pth, output_dir):
+    def _get_spec(row):
+        mzs = np.array([float(m) for m in row["mzs"].split(",")], dtype=np.float32)
+        intensities = np.array([float(i) for i in row["intensities"].split(",")],dtype=np.float32)
+        mzs = float(row['precursor_mz']) - mzs
+        valid_idx = np.where(mzs>1.0)
+        mzs = mzs[valid_idx]
+        intensities = intensities[valid_idx]
+
+        sorted_idx = np.argsort(mzs)
+        mzs = np.concatenate((mzs[sorted_idx], [float(row['precursor_mz'])]))
+        intensities = np.concatenate((intensities[sorted_idx], [1.0]))
+
+        return mzs, intensities
+    
+    spec_data = pd.read_csv(pth, sep='\t')
+    spec_data[['mzs', 'intensities']] = spec_data.apply(lambda row: _get_spec(row), axis=1, result_type='expand')
+    
+    file_pth = os.path.join(output_dir, 'NL_spec.pkl')
+    with open(file_pth, 'wb') as f:
+        pickle.dump(spec_data, f)
+    return file_pth
+
+def generate_cons_spec(pth, output_dir):
+    spec_data = pd.read_csv(pth, sep='\t')
+    data_by_smiles = spec_data[['identifier', 'smiles', 'mzs', 'intensities', 'fold']].groupby('smiles').agg({'identifier':list, 'mzs':lambda x: ','.join(x), 'intensities': lambda x: ','.join(x), 'fold':list})
+    smiles_to_fold  = dict(zip(data_by_smiles.index.tolist(), data_by_smiles['fold'].tolist()))
+
+    consensus_spectra = {}
+    for idx, row in tqdm(data_by_smiles.iterrows(), total=len(data_by_smiles)):
+        mzs = np.array([float(m) for m in row["mzs"].split(",")], dtype=np.float32)
+        intensities = np.array([float(i) for i in row["intensities"].split(",")],dtype=np.float32)
+
+        sorted_idx = np.argsort(mzs)
+        mzs = mzs[sorted_idx]
+        intensities = intensities[sorted_idx]
+        smiles = row.name
+        
+        consensus_spectra[smiles] = {'mzs':mzs, 'intensities':intensities,'precursor_mz': 10000.0,
+                    'fold': smiles_to_fold[smiles][0]}
+    
+    df = pd.DataFrame.from_dict(consensus_spectra, orient='index')
+    df = df.rename_axis('smiles').reset_index()
+
+    return df
+
+
+def generate_cons_spec_formulas(pth, subformula_dir, output_dir=''):
+    # load tsv file
+    spec_data = pd.read_csv(pth, sep='\t')
+
+    # goup spectra by SMILES
+    data_by_smiles = spec_data[['identifier', 'smiles', 'fold', 'precursor_mz', 'formula', 'adduct']].groupby('smiles').agg({'identifier':list, 'fold': list, 'formula': list, 'precursor_mz': "max", 'adduct': list})
+    smiles_to_id = dict(zip(data_by_smiles.index.tolist(), data_by_smiles['identifier'].tolist()))
+    smiles_to_fold  = dict(zip(data_by_smiles.index.tolist(), data_by_smiles['fold'].tolist()))
+    smiles_to_precursorMz = dict(zip(data_by_smiles.index.tolist(), data_by_smiles['precursor_mz'].tolist()))
+    smiles_to_precursorFormula = dict(zip(data_by_smiles.index.tolist(), data_by_smiles['formula'].tolist()))
+    # load subformulas
+    subformulaLoader = data_utils.Subformula_Loader(spectra_view='SpecFormula', dir_path=subformula_dir)
+    id_to_spec = subformulaLoader(spec_data['identifier'].tolist())
+
+    # combine spectra
+    consensus_spectra = {}
+    for smiles, ids in tqdm(smiles_to_id.items(), total=len(data_by_smiles)):
+        cons_spec = collections.defaultdict(list)
+        for id in ids:
+            if id in id_to_spec:
+                for k, v in id_to_spec[id].items():
+                    cons_spec[k].extend(v)
+        cons_spec = pd.DataFrame(cons_spec)
+
+        assert(len(set(smiles_to_fold[smiles]))==1)
+
+        # keep maxed mz and maxed intensity
+        try:
+            cons_spec = cons_spec.groupby('formulas').agg({'formula_mzs': "max", 'formula_intensities': "max"})
+            cons_spec.reset_index(inplace=True)
+        except:
+            d = {
+                'formulas': [smiles_to_precursorFormula[smiles][0]],
+                'formula_mzs': [smiles_to_precursorMz[smiles]],
+                'formula_intensities': [1.0]
+            }
+            cons_spec = pd.DataFrame(d)
+        
+        cons_spec = cons_spec.sort_values(by='formula_mzs').reset_index(drop=True)
+        cons_spec = {'formulas': cons_spec['formulas'].tolist(),
+                    'formula_mzs': cons_spec['formula_mzs'].tolist(),
+                    'formula_intensities': cons_spec['formula_intensities'].tolist(), 
+                    'precursor_mz': smiles_to_precursorMz[smiles],
+                    'fold': smiles_to_fold[smiles][0],
+                    'precursor_formula': smiles_to_precursorFormula[smiles][0]}# formula without adduct...
+
+        consensus_spectra[smiles] = cons_spec
+
+    # save consensus spectra
+    df = pd.DataFrame.from_dict(consensus_spectra, orient='index')
+    df = df.rename_axis('smiles').reset_index()
+
+    return df