main code

mvp/data/__init__.py

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

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,430 @@
+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,
+                **kwargs):
+        
+        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)
+        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]
+            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")
+
+
+            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
+
+class MassSpecDataset_Candidates:
+
+    def __init__(self, 
+                use_formulas: bool,
+                aug_cands_dir_pth: str,
+                aug_cands_size: int,
+                **kwargs):
+        self.aug_cands_size = aug_cands_size
+        self.instance = MassSpecDataset_PeakFormulas(**kwargs, return_mol_freq=False) if use_formulas else JESTR1_MassSpecDataset(**kwargs, return_mol_freq=False)
+
+        with open(aug_cands_dir_pth, 'rb') as f:
+            aug_cands = pickle.load(f)
+
+        if self.use_fp:
+            self.fpgen = AllChem.GetMorganGenerator(radius=5,fpSize=1024)
+
+        self.aug_cands = {}
+        targets = np.array(list(aug_cands.keys()))
+        for smiles, cands in aug_cands.items():
+            # sort candidates by tanimoto similarity
+            cands.sort(key=lambda x: x[1], reverse=True)
+            cands = [c for c in cands if '.' not in c]
+            # assert(len(cands) >0)
+            if len(cands) <=1: # if no candidates, shuffle from target list
+                np.random.shuffle(targets)
+                cands = targets
+            self.aug_cands[smiles] = itertools.cycle(cands)
+
+    def __getattr__(self, name):
+        return self.instance.__getattribute__(name)
+    
+    def __getitem__(self, i):
+        item = self.instance.__getitem__(i,transform_mol=False)
+
+        aug_cands = [next(self.aug_cands[item['mol']]) for _ in range(self.aug_cands_size)]
+        item['aug_cands_fp'] = [self.fpgen.GetFingerprint(Chem.MolFromSmiles(c)).ToList() for c in aug_cands] 
+        item["aug_cands"] = [self.mol_transform(c) for c in aug_cands]
+        item["mol"] = self.mol_transform(item["mol"])
+
+        return item

mvp/data/transforms.py

@@ -0,0 +1,298 @@
+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 SpecMzIntTokenizer(SpecTransform):
+    def __init__(self, max_mz, mz_mean_std=None, mask_precursor=None):
+        self.max_mz = max_mz
+        self.mz_mean_std = mz_mean_std
+    def matchms_transforms(self, spec: matchms.Spectrum):
+        return default_matchms_transforms(spec, mz_to=self.max_mz, n_max_peaks=None)
+    
+    def matchms_to_torch(self, spec: matchms.Spectrum):
+        mzs = spec.peaks.mz
+        intensities = spec.peaks.intensities
+        spec = np.zeros((len(mzs), 2))
+
+        if self.mz_mean_std:
+            mz = (mzs-self.mz_mean_std['mz_mean'])/self.mz_mean_std['mz_std']
+        else:
+            mz = mzs/self.max_mz
+        
+        spec[:, 0] = mz
+        spec[:, 1] = intensities
+
+        return torch.from_numpy(spec.astype(np.float32))
+
+class SpecFormulaMzFeaturizer(SpecTransform):
+    ''' Uses raw mz and intensities '''
+
+    def __init__(
+            self,
+            add_intensities: bool,
+            max_mz: float = 1005,
+            element_list: list = CHEM_ELEMS_SMALL,
+            formula_normalize_vector: Optional[np.array] = None,
+            mz_mean_std: dict[str, float] = None,
+            mask_precursor: bool = False,
+    ) -> None:
+        self.max_mz = max_mz
+        self.elem_to_pos = {e: i for i, e in enumerate(element_list)}
+        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]*)"
+        self.mz_mean_std = mz_mean_std
+        self.add_intensities = add_intensities
+        self.mask_precursor = mask_precursor
+        
+    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'] # mz to formula dict
+
+        peak_idx = np.where(mzs <= self.max_mz)[0]
+        mzs = mzs[peak_idx]
+        intensities = intensities[peak_idx]
+        formulas = [formulas.get(mz, "NA") for mz in mzs[peak_idx]]
+
+        if self.mask_precursor:
+            try:
+                precursor_i = formulas.index(spec.metadata['precursor_formula'])
+                formulas[precursor_i] = 'NA'
+            except:
+                pass
+
+        formulas = self._featurize_formula(formulas)
+        formulas = formulas/self.formula_normalize_vector
+        
+        if self.mz_mean_std:
+            mz = (mzs-self.mz_mean_std['mz_mean'])/self.mz_mean_std['mz_std']
+        else:
+            mz = mzs/self.max_mz
+        
+        if self.add_intensities:
+            spec = np.concatenate((mz.reshape(-1,1), intensities.reshape(-1,1), formulas), axis=1)
+        else:
+            spec = np.concatenate((mz.reshape(-1,1), formulas), axis=1)
+
+        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):
+            if f == "NA":
+                # formula_vector[i] = np.zeros((1, len(self.elem_to_pos)))
+                formula_vector[i] = np.ones((1, len(self.elem_to_pos))) * -1
+
+            else:
+                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
+        return formula_vector
+    
+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,43 @@
+"""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']
+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_FORMULA_VECTOR_NORM = [102.0, 59.0, 25.0, 13.0, 3.0, 6.0, 6.0, 17.0, 4.0, 4.0]
+MSGYM_FORMULA_STANDARD = {
+    'formula_norm':[6.53758314e+00, 6.26973237e+00,
+       8.90610447e-01, 4.73889402e-01, 2.31793513e-02, 3.56956333e-02,
+       2.78056172e-02, 3.28356898e-02, 2.19480328e-03, 1.58458297e-03,
+       2.34802165e-05, 1.71127001e-05, 1.71127001e-04, 1.69800435e-05],
+    'formula_norm': [9.68749281e+00, 7.46795232e+00,
+       1.75427539e+00, 9.81685190e-01, 1.52363430e-01, 2.01197446e-01,
+       1.93046421e-01, 2.65309185e-01, 4.82433244e-02, 5.23009413e-02,
+       4.84558202e-03, 4.13671455e-03, 1.51218609e-02, 5.02040474e-03],
+    'formula_max':[102.        ,  59.        ,
+        25.        ,  13.        ,   3.        ,   6.        ,
+         6.        ,  17.        ,   4.        ,   4.        ,
+         1.        ,   1.        ,   5.        ,   2.        ],
+    'formula_min' :[0.        , 0.        , 0.        , 0.        ,
+       0.        , 0.        , 0.        , 0.        , 0.        ,
+       0.        , 0.        , 0.        , 0.        , 0.        ,
+       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,799 @@
+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 torch.nn.utils.rnn import pad_sequence
+
+from mvp.utils.loss import contrastive_loss, cand_spec_sim_loss, fp_loss, cons_spec_loss, filip_loss_with_mask
+import mvp.utils.models as model_utils
+from mvp.utils.general import pad_graph_nodes, filip_similarity_batch
+
+from mvp.models.encoders import CrossAttention
+import torch.nn.functional as F
+
+from torch_geometric.nn import global_mean_pool
+
+class ContrastiveModel(RetrievalMassSpecGymModel):
+    def __init__(
+        self,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.save_hyperparameters()
+
+        if 'use_fp' not in self.hparams:
+            self.hparams.use_fp = False
+        if 'use_fp' not in self.hparams:
+            self.hparams.use_fp = False
+        if 'use_NL_spec' not in self.hparams:
+            self.hparams.use_NL_spec = 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)
+        
+        # setup loss strategy
+        if self.hparams.model == 'contrastive':
+            self._loss_setup()
+        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 _loss_setup(self):
+        self.loss_wts = {}
+        self.loss_updates = {}
+
+
+        for p, loss in zip(['use_contr','pred_fp', 'use_cons_spec', 'aug_cands'], ['contr_wt','fp_wt','cons_spec_wt' ,'aug_cands_wt']):
+            if p not in self.hparams:
+                self.hparams[p] = False
+            if self.hparams[p]:
+                if self.hparams.loss_strategy == 'linear':
+                    start_wt = self.hparams[loss+'_update']['start']
+                    end_wt = self.hparams[loss+'_update']['end']
+                    change = (end_wt - start_wt)/self.hparams.max_epochs
+                    self.loss_updates[loss] = change
+                    self.loss_wts[loss] = start_wt
+                elif self.hparams.loss_strategy == 'manual':
+                    self.loss_updates[loss] = self.hparams[loss+'_update']
+                    self.loss_wts[loss] = self.hparams[loss]
+                else: 
+                    self.loss_wts[loss] = self.hparams[loss]
+                    
+    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):
+        # 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):
+        # 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 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
+        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)
+
+        self.df_test.to_pickle(self.df_test_path)
+
+class FilipContrastive(ContrastiveModel):
+    def __init__(self,
+                 **kwargs):
+        
+        super().__init__(**kwargs)
+
+    def compute_loss(self, batch: dict, spec_enc, mol_enc, spec_mask, mol_mask):
+        losses = {}
+
+        loss = filip_loss_with_mask(spec_enc, mol_enc, spec_mask, mol_mask, self.hparams.contr_temp)
+
+        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)
+
+        # pad nodes to max_n_nodes in batch (Spectra are already padded)
+        mol_enc, mol_mask = pad_graph_nodes(mol_enc, batch['mol_n_nodes'])
+        spec_mask = ~torch.all((spec_enc == -5), dim=-1)
+
+        if stage == Stage.TEST:
+            return dict(spec_enc=spec_enc, mol_enc=mol_enc, spec_mask=spec_mask, mol_mask=mol_mask)
+
+        # Calculate loss
+        losses = self.compute_loss(batch, spec_enc, mol_enc, spec_mask, mol_mask)
+
+        return losses
+    
+    def test_step(self, batch):
+        # 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']
+        spec_mask = outputs['spec_mask']
+        mol_mask = outputs['mol_mask']
+
+        # Calculate scores
+        indexes = utils.batch_ptr_to_batch_idx(batch_ptr)
+        
+        scores = filip_similarity_batch(spec_enc, mol_enc, spec_mask, mol_mask)
+        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)
+    
+class MultiViewFineTuning(MultiViewContrastive):
+    def __init__(self,
+                 **kwargs):
+        super().__init__(**kwargs)
+
+        # load preptrained spec, mol, fp encoders
+        checkpoint = torch.load(self.hparams.partial_checkpoint)
+        state_dict = state_dict = {k[len("spec_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("spec_enc_model")}
+        self.spec_enc_model.load_state_dict(state_dict) # trained on consensus spectra
+
+        state_dict = state_dict = {k[len("mol_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("mol_enc_model")}
+        self.mol_enc_model.load_state_dict(state_dict)
+
+        state_dict = state_dict = {k[len("fp_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("fp_enc_model")}
+        self.fp_enc_model.load_state_dict(state_dict)
+
+        self.encoding_views = ['spec_enc', 'mol_enc', 'fp_enc']
+        self.loss_fn = nn.BCELoss()
+
+        # freeze encoders
+        for param in self.mol_enc_model.parameters():
+            param.requires_grad = False
+        for param in self.spec_enc_model.parameters():
+            param.requires_grad = False
+        for param in self.fp_enc_model.parameters():
+            param.requires_grad = False
+        for param in self.cons_spec_enc_model.parameters():
+            param.requires_grad = False
+
+        # n_views = 2
+        # if self.hparams.use_fp:
+        #     n_views+=1
+
+        # in_dim = self.hparams.final_embedding_dim*n_views
+        in_dim = self.hparams.final_embedding_dim *2 + 2
+
+        self.classifier_model = nn.Sequential(
+            nn.Linear(in_dim, 512),
+            nn.ReLU(),
+            nn.BatchNorm1d(512),
+            nn.Dropout(0.3),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.BatchNorm1d(256),
+            nn.Dropout(0.3),
+            nn.Linear(256, 1),
+            nn.Sigmoid()
+        )
+        self.noise_std = 0.01
+
+    def _add_noise(self, x):
+        noise = torch.randn_like(x) * self.noise_std
+        return x + noise
+
+    def forward(self, batch, stage):
+
+        matching_views = super().forward(batch, stage)
+        # matching_enc = torch.concat((matching_views['spec_enc'], matching_views['mol_enc'], matching_views['fp_enc']), dim=-1)
+        # enc1 = matching_views['spec_enc'] - matching_views['mol_enc']
+        # enc2 = matching_views['spec_enc'] - matching_views['fp_enc']
+        # matching_enc = torch.concat((enc1, enc2), dim=-1)
+        view1 = matching_views['spec_enc']
+        view2 = matching_views['mol_enc']
+        view3 = matching_views['fp_enc']
+
+        if stage == Stage.TRAIN:
+            view1, view2, view3 = map(self._add_noise, (view1, view2, view3))
+
+        pairwise_diffs = torch.cat([
+            torch.abs(view1 - view2),
+            torch.abs(view1 - view3),
+        ], dim=-1)
+
+        pairwise_sims = torch.cat([
+            (view1 * view2).sum(dim=-1, keepdim=True),
+            (view1 * view3).sum(dim=-1, keepdim=True),
+        ], dim=-1)
+
+        matching_enc = torch.cat([pairwise_diffs, pairwise_sims], dim=-1)
+        matching_scores = self.classifier_model(matching_enc)
+
+        if stage == Stage.TEST:
+            return dict(matching_scores = matching_scores)
+        
+        view1 = view1.repeat_interleave(self.hparams.aug_cands_size, dim=0)
+        view2 = self.mol_enc_model(batch['aug_cands'])
+        view3= self.fp_enc_model(batch['aug_cands_fp'])
+        if stage == Stage.TRAIN:
+            view1, view2, view3 = map(self._add_noise, (view1, view2, view3))
+
+        pairwise_diffs = torch.cat([
+            torch.abs(view1 - view2),
+            torch.abs(view1 - view3),
+        ], dim=-1)
+
+        pairwise_sims = torch.cat([
+            (view1 * view2).sum(dim=-1, keepdim=True),
+            (view1 * view3).sum(dim=-1, keepdim=True),
+        ], dim=-1)
+
+        nonmatching_enc = torch.cat([pairwise_diffs, pairwise_sims], dim=-1)
+
+        nonmatching_scores = self.classifier_model(nonmatching_enc)
+
+        return dict(matching_scores=matching_scores, nonmatching_scores=nonmatching_scores)
+    
+    def compute_loss(self, matching_scores, nonmatching_scores):
+        
+        matching_loss = self.loss_fn(matching_scores, torch.ones_like(matching_scores).to(matching_scores.device))
+        nonmatching_loss = self.loss_fn(nonmatching_scores, torch.zeros_like(nonmatching_scores).to(nonmatching_scores.device))
+
+        loss = matching_loss + (1/self.hparams.aug_cands_size)*nonmatching_loss
+
+        return dict(loss=loss)
+    
+    def step(
+        self, batch: dict, stage= Stage.NONE):
+        
+        output = self.forward(batch, stage)
+
+        if stage == Stage.TEST:
+            return output
+        
+        # Calculate loss
+        losses = self.compute_loss(output['matching_scores'], output['nonmatching_scores'])
+
+        return losses
+    
+    def test_step(self, batch):
+        # 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 = outputs['matching_scores']
+
+        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_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
+        )
+
+    def on_test_batch_end(self, outputs, batch: dict, batch_idx: int, stage: Stage = Stage.TEST) -> None:
+        ContrastiveModel.on_test_batch_end(self, outputs, batch, batch_idx, stage)
+
+    def on_test_epoch_end(self):
+        self.df_test = pd.DataFrame.from_dict(self.result_dct, orient='index').reset_index().rename(columns={'index': 'identifier'})
+        # self.df_test.to_csv(self.hparams.resutl)
+        print(self.df_test_path)
+        self.df_test.to_pickle(self.df_test_path)
+        # ContrastiveModel.on_test_epoch_end(self)
+
+    def get_checkpoint_monitors(self) -> T.List[dict]:
+        monitors = [
+            {"monitor": f"{Stage.VAL.to_pref()}loss", "mode": "min", "early_stopping": True}
+        ]
+        return monitors
+    def configure_optimizers(self):
+        return torch.optim.Adam(
+            self.classifier_model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay
+        )
+
+class IndSpecEncoder(ContrastiveModel):
+    """ Trains a spectra encoder that maps to a pretrained spec encoder"""
+    def __init__(
+            self,
+            **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        # initialize ind_spec_encoder and loss
+        self.ind_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)
+
+        # load preptrained spec and mol encoders
+        checkpoint = torch.load(self.hparams.partial_checkpoint)
+        state_dict = state_dict = {k[len("spec_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("spec_enc_model")}
+        self.spec_enc_model.load_state_dict(state_dict) # trained on consensus spectra
+
+        state_dict = state_dict = {k[len("mol_enc_model."):]: v for k, v in checkpoint['state_dict'].items() if k.startswith("mol_enc_model")}
+        self.mol_enc_model.load_state_dict(state_dict)
+
+        # freeze cons spec and mol encoders
+        for param in self.mol_enc_model.parameters():
+            param.requires_grad = False
+        for param in self.spec_enc_model.parameters():
+            param.requires_grad = False
+
+    def forward(self, batch, stage):
+
+        spec = batch[self.spec_view]
+        n_peaks = batch['n_peaks']
+        spec_enc = self.ind_spec_enc_model(spec, n_peaks)
+
+        return spec_enc
+    
+    def compute_loss(self, spec_enc, cons_spec_enc):
+        loss = self.cons_loss(spec_enc, cons_spec_enc)
+        return dict(loss=loss)
+
+    def step(self, batch: dict, stage=Stage.NONE):
+        self.spec_enc_model.eval()
+        self.mol_enc_model.eval()
+
+        spec_enc = self.forward(batch, stage)
+
+        if stage == Stage.TEST:
+            mol_enc = self.mol_enc_model(batch['cand'])
+            return dict(spec_enc=spec_enc, mol_enc=mol_enc)
+        
+        cons_spec_enc = self.spec_enc_model(batch['cons_spec'], batch['cons_n_peaks'])
+
+        losses = self.compute_loss(spec_enc, cons_spec_enc)
+
+        return losses
+    
+    
+    def configure_optimizers(self):
+        return torch.optim.Adam(
+            self.ind_spec_enc_model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay
+        )
+    def get_checkpoint_monitors(self) -> T.List[dict]:
+        monitors = [
+            {"monitor": f"{Stage.VAL.to_pref()}loss", "mode": "min", "early_stopping": True}
+        ]
+        return monitors
+    
+class CrossAttenContrastive(ContrastiveModel):
+    def __init__(
+            self,
+            **kwargs
+    ):
+        super(CrossAttenContrastive, self).__init__(**kwargs)
+        self.specMolCrossAttentionModel = CrossAttention(self.hparams.formula_dims[-1], self.hparams.gnn_channels[-1], self.hparams.cross_attn_heads, dim_out=self.hparams.final_embedding_dim, dropout=0.3)
+        self.molSpecCrossAttentionModel = CrossAttention(self.hparams.gnn_channels[-1], self.hparams.formula_dims[-1], self.hparams.cross_attn_heads, dim_out=self.hparams.final_embedding_dim, dropout=0.3)
+
+    def forward(self, batch, stage) -> tuple[torch.Tensor, torch.Tensor]:
+        # Unpack inputs
+        spec = batch[self.spec_view]
+        spec_n_forms = batch['n_peaks']
+        g = batch['cand'] if stage == Stage.TEST else batch['mol']
+        g_n_nodes = batch['mol_n_nodes']
+
+        # encode peaks and nodes
+        spec_enc = self.spec_enc_model(spec)
+        mol_enc = self.mol_enc_model(g)
+    
+        # pad mol_enc and spec_enc to have the same length
+        max_nodes = max(g_n_nodes)
+        max_forms = max(spec_n_forms)
+
+        if max_forms > max_nodes: ## pad mol_enc
+            mol_enc = torch.cat((mol_enc, torch.rand(max_forms, self.hparams.gnn_channels[-1]).to(spec.device)))
+            mol_enc = torch.split(mol_enc, g_n_nodes+[max_forms])
+            mol_enc = pad_sequence(mol_enc, batch_first=True, padding_value=-5)[:-1,:,:]
+
+        elif max_nodes > max_forms: ## pad spec_enc
+            dim_diff = max_nodes - max_forms
+            spec_enc = F.pad(spec_enc, (0,0,0,dim_diff, 0,0), value=-5)
+            mol_enc = torch.split(mol_enc, g_n_nodes)
+            mol_enc = pad_sequence(mol_enc, batch_first=True, padding_value=-5)
+        else:
+            mol_enc = torch.split(mol_enc, g_n_nodes)
+            mol_enc = pad_sequence(mol_enc, batch_first=True, padding_value=-5)
+
+        spec_pad = torch.all((spec_enc == -5), -1)
+        mol_pad = torch.all((mol_enc == -5), -1)
+
+        # cross attention
+        tmp_spec_enc = spec_enc * 1.0
+        spec_enc = self.specMolCrossAttentionModel(spec_enc, mol_enc, mol_enc, mask=mol_pad)
+        mol_enc = self.molSpecCrossAttentionModel(mol_enc, tmp_spec_enc, tmp_spec_enc, mask=spec_pad)
+
+        # pool
+        spec_indecies = torch.tensor([i for i, count in enumerate(spec_n_forms) for _ in range(count)]).to(spec_enc.device)
+        mol_indecies = torch.tensor([i for i, count in enumerate(g_n_nodes) for _ in range(count)]).to(mol_enc.device)
+
+        spec_enc = spec_enc[~spec_pad].reshape(-1, spec_enc.shape[-1])
+        mol_enc = mol_enc[~mol_pad].reshape(-1, mol_enc.shape[-1])
+
+        spec_enc = global_mean_pool(spec_enc, spec_indecies)
+        mol_enc = global_mean_pool(mol_enc, mol_indecies)
+
+        return spec_enc, mol_enc

mvp/models/encoders.py

@@ -0,0 +1,74 @@
+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
+
+class CrossAttention(nn.Module):
+    def __init__(self, embed_dim_q, embed_dim_kv, num_heads, dim_out,  dropout=0.0):
+        """
+        Args:
+            embed_dim_q (int): Dimension of query embeddings.
+            embed_dim_kv (int): Dimension of key/value embeddings.
+            num_heads (int): Number of attention heads.
+            dropout (float): Dropout probability for attention weights.
+        """
+        super(CrossAttention, self).__init__()
+       
+        # Ensure the embedding dimensions are divisible by the number of heads
+        assert embed_dim_q % num_heads == 0, "embed_dim_q must be divisible by num_heads"
+        assert embed_dim_kv % num_heads == 0, "embed_dim_kv must be divisible by num_heads"
+
+        self.query_proj = nn.Linear(embed_dim_q, embed_dim_q)
+        self.key_proj = nn.Linear(embed_dim_kv, embed_dim_q)  # Match dimensions with queries
+        self.value_proj = nn.Linear(embed_dim_kv, embed_dim_q)
+
+        self.attention = nn.MultiheadAttention(embed_dim=embed_dim_q, num_heads=num_heads, dropout=dropout, batch_first=True)
+        self.out_proj = nn.Linear(embed_dim_q, dim_out)
+
+    def forward(self, queries, keys, values, mask=None):
+        """
+        Args:
+            queries (Tensor): Shape (batch_size, len_q, embed_dim_q)
+            keys (Tensor): Shape (batch_size, len_k, embed_dim_kv)
+            values (Tensor): Shape (batch_size, len_v, embed_dim_kv)
+            mask (Tensor, optional): Shape (batch_size, len_q, len_k), 1 for valid positions and 0 for masked.
+
+        Returns:
+            Tensor: Shape (batch_size, len_q, embed_dim_q)
+        """
+        # Project inputs to the required dimensions
+        queries = self.query_proj(queries)  # (batch_size, len_q, embed_dim_q)
+        keys = self.key_proj(keys)          # (batch_size, len_k, embed_dim_q)
+        values = self.value_proj(values)    # (batch_size, len_v, embed_dim_q)
+
+        # Compute attention
+        attn_output, _ = self.attention(queries, keys, values, key_padding_mask=mask)
+
+        # Apply output projection
+        output = self.out_proj(attn_output)
+        return output

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,182 @@
+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 SpecMzIntTokenTransformer(nn.Module):
+    def __init__(self, args):
+        super(SpecMzIntTokenTransformer, self).__init__()
+        in_dim = 2
+        self.tokenEnc = MLP(in_dim, args.hidden_dims, dropout=args.peak_dropout)
+
+        self.returnEmb = False
+        if args.model in ('crossAttenContrastive', 'filipContrastive'):
+            self.returnEmb = True
+            assert(args.use_cls == False)
+
+        self.use_cls = args.use_cls
+        if self.use_cls:
+            self.cls_embed = torch.nn.Embedding(1,args.hidden_dims[-1])
+        encoder_layer = nn.TransformerEncoderLayer(d_model=args.hidden_dims[-1], nhead=2, batch_first=True)
+        self.tokenTransformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
+
+        self.specEncoder = nn.Sequential(nn.Linear(args.hidden_dims[-1], args.final_embedding_dim), nn.Dropout(args.fc_dropout))
+    
+    def forward(self, spec, n_peaks=None):
+        h = self.tokenEnc(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.tokenTransformer(h, src_key_padding_mask=pad)
+            h = h[:,0,:]
+        else:
+            
+            # mean
+            h = self.tokenTransformer(h, src_key_padding_mask=pad)
+
+            if self.returnEmb:
+                # repad h
+                h[pad] = -5
+                return h
+            n_peaks_indices = torch.tensor([i for i, count in enumerate(n_peaks) for _ in range(count)]).to(spec.device)
+            h = h[~pad].reshape(-1, h.shape[-1])
+            h = global_mean_pool(h, n_peaks_indices)
+
+        h = self.specEncoder(h)
+        return h
+
+    
+class SpecFormulaEncMLP(nn.Module):
+    def __init__(self, args, out_dim=None):
+        super(SpecFormulaEncMLP, self).__init__()
+        in_dim = len(args.element_list)
+        if args.add_intensities:
+            in_dim+=1
+        if args.spectra_view == "SpecFormulaMz": #mz
+            in_dim+=1 
+
+        self.formulaEnc = MLP(in_dim, args.formula_dims, dropout=args.formula_dropout)
+
+        if not out_dim:
+            out_dim = args.final_embedding_dim
+        self.mz_fc1 = nn.Linear(args.formula_dims[-1], out_dim)
+        self.dropout = nn.Dropout(args.fc_dropout)
+
+    def forward(self, spec, n_peaks):
+        h = self.formulaEnc(spec)
+        h = torch.sum(h, axis=1)
+
+        h = self.mz_fc1(h)
+        h = self.dropout(h)
+        return h
+    
+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
+        if args.model in ('crossAttenContrastive', 'filipContrastive'):
+            self.returnEmb = True
+            assert(args.use_cls == 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
+    
+class SpecFormula_mz_Encoder(nn.Module):
+    '''
+    Encodes formula and mz_int
+    '''
+
+    def __init__(self, args):
+        
+        super(SpecFormula_mz_Encoder, self).__init__()
+        
+        self.formula_encoder = SpecFormulaTransformer(args, out_dim=args.final_embedding_dim//4)
+        self.mz_encoder = SpecEncMLP_BIN(args, out_dim=args.final_embedding_dim//4)
+
+        self.fc = nn.Sequential(nn.Linear(args.final_embedding_dim //2, args.final_embedding_dim), nn.ReLU(),
+                            )
+
+    def forward(self, formulas, binned_mzs):
+        h_formula = self.formula_encoder(formulas)
+        h_bin = self.mz_encoder(binned_mzs)
+
+        h_spec = torch.concat((h_formula, h_bin), axis=1)
+        h = self.fc(h_spec)
+
+        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_quick_test'
+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: /r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_mass.json # "../data/MassSpecGym/data/molecules/MassSpecGym_retrieval_candidates_formula.json"
+dataset_pth: /r/hassounlab/spectra_data/msgym/MassSpecGym.tsv #/data/yzhouc01/spectra_data/combined_msgym_nist23_multiplex.tsv # /r/hassounlab/spectra_data/msgym/MassSpecGym.tsv # "../data/MassSpecGym/data/sample_data.tsv" 
+subformula_dir_pth: /data/yzhouc01/MVP/data/MassSpecGym/data/subformulae_default #/data/yzhouc01/spectra_data/subformulae #"../data/MassSpecGym/data/subformulae_default"
+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 #SpecMzIntTokens #SpecFormula 
+#  1. Binner
+max_mz: 1000
+bin_width: 1
+mask_peak_ratio: 0.00
+
+# 2. SpecFormula
+element_list: ['H', 'C',  'O', 'N', 'P', 'S', 'Cl', 'F', 'Br', 'I', 'B', 'As', 'Si', 'Se']
+add_intensities: True
+mask_precursor: False
+
+# - Molecule
+molecule_view: "MolGraph"
+atom_feature: 'full'
+bond_feature: 'full'
+
+
+############################## Views ##############################
+# contrastive
+use_contr: False
+contr_wt: 1
+contr_wt_update: {} 
+
+# consensus spectra
+use_cons_spec: False
+cons_spec_wt: 3
+cons_spec_wt_update: {} 
+cons_loss_type: 'l2' # cosine, l2
+
+# fp prediction/usage
+pred_fp: False
+use_fp: False
+fp_loss_type: 'cosine' #cosine, bce
+fp_wt: 3
+fp_wt_update: {} 
+fp_size: 1024
+fp_radius: 5
+fp_dropout: 0.4
+
+# candidates
+aug_cands: False
+aug_cands_wt: 0.1
+aug_cands_update: {} 
+aug_cands_size: 3
+
+# neutral loss
+use_NL: False
+
+
+############################## Task and model ##############################
+task: 'retrieval'
+spec_enc: Transformer_Formula # Transformer_MzInt #Transformer_Formula 
+mol_enc: "GNN"
+model:  filipContrastive # "MultiviewContrastive"
+contr_views: [['spec_enc', 'mol_enc']] #[['spec_enc', 'mol_enc'], ['spec_enc', 'NL_spec_enc'], ['mol_enc', 'NL_spec_enc']] #[['spec_enc', 'mol_enc'], ['mol_enc', 'cons_spec_enc'], ['cons_spec_enc', 'spec_enc'], ['fp_enc', 'mol_enc'], ['fp_enc', 'spec_enc'], ['fp_enc', 'cons_spec_enc']]
+log_only_loss_at_stages: []
+df_test_path: ""
+
+# - Spectra encoder 
+final_embedding_dim: 512
+fc_dropout: 0.4
+
+# - Spectra Token encoder
+hidden_dims: [64, 128]
+peak_dropout: 0.2
+
+# - Formula-based spec encoders
+formula_dropout: 0.2
+formula_dims: [64, 128, 256]
+cross_attn_heads: 2
+use_cls: False
+
+# -- GAT params
+attn_heads: [12,12,12]
+
+# - Molecule encoder (GNN)
+gnn_channels: [64,128,256] 
+gnn_type: "gcn"
+num_gnn_layers: 3
+gnn_hidden_dim: 512
+gnn_dropout: 0.3

mvp/params_jestr.yaml

@@ -0,0 +1,122 @@
+
+# Experiment setup 
+job_key: ''
+run_name: 'combined_d_1024dim_100bs'
+run_details: ""
+project_name: ''
+wandb_entity_name: 'mass-spec-ml'
+no_wandb: True
+seed: 3
+debug: False
+checkpoint_pth:
+
+# Training setup
+max_epochs: 2000 
+accelerator: 'gpu'
+devices: [1]
+log_every_n_steps: 250
+val_check_interval: 1.0
+
+# Data paths
+candidates_pth: "/r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_mass.json"
+dataset_pth: '/r/hassounlab/spectra_data/msgym/MassSpecGym.tsv' # '/r/hassounlab/spectra_data/msgym/MassSpecGym.tsv' #"/data/yzhouc01/spectra_data/combined_msgym_nist23_multiplex.tsv" 
+subformula_dir_pth: ""
+split_pth: 
+fp_dir_pth: ''
+cons_spec_dir_pth: 
+NL_spec_dir_pth: ""
+partial_checkpoint: ""
+
+# General hyperparameters
+batch_size: 100
+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: 1024
+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: 1024
+gnn_dropout: 0.3

mvp/run.sh

@@ -0,0 +1,12 @@
+# 1. preprocess data (subformula labels should be obtained through MIST)
+# python data_preprocess.py --spec_type formSpec --dataset_pth ../data/sample/data.tsv --candidates_pth  ../data/sample/candidates_mass.json --subformula_dir_pth ../data/sample/subformulae_default/ --output_dir ../data/sample/
+
+# 2. train model on msgym
+# python train.py --param_pth params_formSpec.yaml
+
+# 3. test model on msgym
+# python train.py --param_pth params_binnedSpec.yaml
+
+# python train.py
+python test.py
+python test.py --candidates_pth /r/hassounlab/spectra_data/msgym/molecules/MassSpecGym_retrieval_candidates_formula.json

mvp/subformula_assign/assign_subformulae.py

@@ -0,0 +1,214 @@
+""" assign_subformulae.py
+
+Copied from https://github.com/samgoldman97/mist/blob/main_v2/src/mist/subformulae/assign_subformulae.py
+
+Given a set of spectra and candidates from a labels file, assign subformulae and save to JSON files.
+
+"""
+
+from pathlib import Path
+import argparse
+from functools import partial
+import numpy as np
+import pandas as pd
+import json
+
+from tqdm import tqdm
+import utils
+
+
+def get_args():
+    """get args"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--feature-id",
+        default="ID",
+        help="ID key in mgf input"
+    )
+    parser.add_argument(
+        "--spec-files",
+        default="data/paired_spectra/canopus_train/spec_files/",
+        help="Spec files; either MGF or directory.",
+    )
+    parser.add_argument("--output-dir", default=None,
+                        help="Name of output dir.")
+    parser.add_argument(
+        "--labels-file",
+        default="data/paired_spectra/canopus_train/labels.tsv",
+        help="Labels file",
+    )
+    parser.add_argument(
+        "--debug", action="store_true", default=False, help="Debug flag."
+    )
+    parser.add_argument(
+        "--mass-diff-type",
+        default="ppm",
+        type=str,
+        help="Type of mass difference - absolute differece (abs) or relative difference (ppm).",
+    )
+    parser.add_argument(
+        "--mass-diff-thresh",
+        action="store",
+        default=20,
+        type=float,
+        help="Threshold of mass difference.",
+    )
+    parser.add_argument(
+        "--inten-thresh",
+        action="store",
+        default=0.001,
+        type=float,
+        help="Threshold of MS2 subpeak intensity (normalized to 1).",
+    )
+    parser.add_argument(
+        "--max-formulae",
+        action="store",
+        default=50,
+        type=int,
+        help="Max number of peaks to keep",
+    )
+    parser.add_argument(
+        "--num-workers", action="store", default=32, type=int, help="num workers"
+    )
+    return parser.parse_args()
+
+
+def process_spec_file(spec_name: str, spec_files: str, max_inten=0.001, max_peaks=60):
+    """_summary_
+
+    Args:
+        spec_name (str): _description_
+        spec_files (str): _description_
+        max_inten (float, optional): _description_. Defaults to 0.001.
+        max_peaks (int, optional): _description_. Defaults to 60.
+
+    Returns:
+        _type_: _description_
+    """
+    spec_file = Path(spec_files) / f"{spec_name}.ms"
+
+    meta, tuples = utils.parse_spectra(spec_file)
+    spec = utils.process_spec_file(meta, tuples)
+    return spec_name, spec
+
+
+def assign_subforms(spec_files, labels_file,
+                    mass_diff_thresh: int = 20,
+                    mass_diff_type: str = "ppm",
+                    inten_thresh: float = 0.001,
+                    output_dir=None,
+                    num_workers: int = 32,
+                    feature_id="ID",
+                    max_formulae: int = 50,
+                    debug=False):
+    """_summary_
+
+    Args:
+        spec_files (_type_): _description_
+        labels_file (_type_): _description_
+        mass_diff_thresh (int, optional): _description_. Defaults to 20.
+        mass_diff_type (str, optional): _description_. Defaults to "ppm".
+        inten_thresh (float, optional): _description_. Defaults to 0.001.
+        output_dir (_type_, optional): _description_. Defaults to None.
+        num_workers (int, optional): _description_. Defaults to 32.
+        feature_id (str, optional): _description_. Defaults to "ID".
+        max_formulae (int, optional): _description_. Defaults to 50.
+        debug (bool, optional): _description_. Defaults to False.
+
+    Raises:
+        ValueError: _description_
+    """
+    spec_files = Path(spec_files)
+    label_path = Path(labels_file)
+
+    # Read in labels
+    labels_df = pd.read_csv(label_path, sep="\t").astype(str)
+    if spec_files.suffix == ".tsv": # YZC msgym-like data
+        labels_df.rename(columns={'identifier': 'spec',
+            'adduct': 'ionization'}, inplace=True)
+
+    if debug:
+        labels_df = labels_df[:50]
+
+    # Define output directory name
+    output_dir = Path(output_dir)
+    if output_dir is None:
+        subform_dir = label_path.parent / "subformulae"
+        output_dir_name = f"subform_{max_formulae}"
+        output_dir = subform_dir / output_dir_name
+
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    if spec_files.suffix == ".mgf":
+        # Input specs
+        parsed_specs = utils.parse_spectra_mgf(spec_files)
+        input_specs = [utils.process_spec_file(*i) for i in parsed_specs]
+        spec_names = [i[0][feature_id] for i in parsed_specs]
+        input_specs = list(zip(spec_names, input_specs))
+    elif spec_files.is_dir():
+        spec_fn_lst = labels_df["spec"].to_list()
+        proc_spec_full = partial(
+            process_spec_file,
+            spec_files=spec_files,
+            max_inten=inten_thresh,
+            max_peaks=max_formulae,
+        )
+        # input_specs = [proc_spec_full(i) for i in tqdm(spec_fn_lst)]
+        input_specs = utils.chunked_parallel(
+            spec_fn_lst, proc_spec_full, chunks=100, max_cpu=max(num_workers, 1)
+        )
+
+    elif spec_files.suffix == '.tsv':
+        parsed_specs = utils.parse_spectra_msgym(labels_df)
+        input_specs = [utils.process_spec_file(*i) for i in parsed_specs]
+        spec_names = [i[0][feature_id] for i in parsed_specs]
+        input_specs = list(zip(spec_names, input_specs))
+    else:
+        raise ValueError(f"Spec files arg {spec_files} is not a dir or mgf")
+    
+
+    # input_specs contains a list of tuples (spec, subpeak tuple array)
+    input_specs_dict = {tup[0]: tup[1] for tup in input_specs}
+    export_dicts, spec_names = [], []
+    for _, row in labels_df.iterrows():
+        spec = str(row["spec"])
+        new_entry = {
+            "spec": input_specs_dict[spec],
+            "form": row["formula"],
+            "mass_diff_type": mass_diff_type,
+            "spec_name": spec,
+            "mass_diff_thresh": mass_diff_thresh,
+            "ion_type": row["ionization"],
+        }
+        spec_names.append(spec)
+        export_dicts.append(new_entry)
+
+    # Build dicts
+    print(f"There are {len(export_dicts)} spec-cand pairs this spec files")
+    def export_wrapper(x): return utils.get_output_dict(**x)
+    if debug:
+        output_dict_lst = [export_wrapper(i) for i in export_dicts[:10]]
+    else:
+        output_dict_lst = utils.chunked_parallel(
+            export_dicts, export_wrapper, chunks=100, max_cpu=max(num_workers, 1)
+        )
+    assert len(export_dicts) == len(output_dict_lst)
+
+    # Write all output jsons to files
+    for output_dict, spec_name in tqdm(zip(output_dict_lst, spec_names)):
+        with open(output_dir / f"{spec_name}.json", "w") as f:
+            json.dump(output_dict, f, indent=4)
+            f.close()
+
+if __name__ == "__main__":
+    args = get_args()
+    assign_subforms(spec_files=args.spec_files, 
+                    labels_file=args.labels_file,
+                    mass_diff_thresh=args.mass_diff_thresh,
+                    mass_diff_type=args.mass_diff_type,
+                    inten_thresh=args.inten_thresh,
+                    output_dir=args.output_dir,
+                    num_workers=args.num_workers,
+                    feature_id=args.feature_id,
+                    max_formulae=args.max_formulae,
+                    debug=args.debug)

mvp/subformula_assign/run.sh

@@ -0,0 +1,6 @@
+SPEC_FILES="/data/yzhouc01/spectra_data/combined_msgym_nist23_multiplex.tsv"
+OUTPUT_DIR="/data/yzhouc01/spectra_data/subformulae"
+MAX_FORMULAE=60
+LABELS_FILE="/data/yzhouc01/spectra_data/combined_msgym_nist23_multiplex.tsv"
+
+python assign_subformulae.py --spec-files $SPEC_FILES --output-dir $OUTPUT_DIR --max-formulae $MAX_FORMULAE --labels-file $LABELS_FILE 

mvp/subformula_assign/utils/__init__.py

@@ -0,0 +1,5 @@
+
+from .parse_utils import *
+from .chem_utils import *
+from .parallel_utils import *
+from .spectra_utils import *

mvp/subformula_assign/utils/chem_utils.py

@@ -0,0 +1,612 @@
+"""chem_utils.py"""
+
+import re
+import numpy as np
+import pandas as pd
+import json
+from functools import reduce
+from collections import defaultdict
+
+import torch
+from rdkit import Chem
+from rdkit.Chem import Atom
+from rdkit.Chem.rdMolDescriptors import CalcMolFormula
+from rdkit.Chem.Descriptors import ExactMolWt
+from rdkit.Chem.MolStandardize import rdMolStandardize
+
+P_TBL = Chem.GetPeriodicTable()
+
+ROUND_FACTOR = 4
+
+ELECTRON_MASS = 0.00054858
+CHEM_FORMULA_SIZE = "([A-Z][a-z]*)([0-9]*)"
+
+VALID_ELEMENTS = [
+    "C",
+    "H",
+    "As",
+    "B",
+    "Br",
+    "Cl",
+    "Co",
+    "F",
+    "Fe",
+    "I",
+    "K",
+    "N",
+    "Na",
+    "O",
+    "P",
+    "S",
+    "Se",
+    "Si",
+]
+VALID_ATOM_NUM = [Atom(i).GetAtomicNum() for i in VALID_ELEMENTS]
+
+
+CHEM_ELEMENT_NUM = len(VALID_ELEMENTS)
+
+ATOM_NUM_TO_ONEHOT = torch.zeros((max(VALID_ATOM_NUM) + 1, CHEM_ELEMENT_NUM))
+
+# Convert to onehot
+ATOM_NUM_TO_ONEHOT[VALID_ATOM_NUM, torch.arange(CHEM_ELEMENT_NUM)] = 1
+
+VALID_MONO_MASSES = np.array(
+    [P_TBL.GetMostCommonIsotopeMass(i) for i in VALID_ELEMENTS]
+)
+CHEM_MASSES = VALID_MONO_MASSES[:, None]
+
+ELEMENT_VECTORS = np.eye(len(VALID_ELEMENTS))
+ELEMENT_VECTORS_MASS = np.hstack([ELEMENT_VECTORS, CHEM_MASSES])
+ELEMENT_TO_MASS = dict(zip(VALID_ELEMENTS, CHEM_MASSES.squeeze()))
+
+ELEMENT_DIM_MASS = len(ELEMENT_VECTORS_MASS[0])
+ELEMENT_DIM = len(ELEMENT_VECTORS[0])
+
+# Reasonable normalization vector for elements
+# Estimated by max counts (+ 1 when zero)
+NORM_VEC = np.array([81, 158, 2, 1, 3, 10, 1, 17, 1, 6, 1, 19, 2, 34, 6, 6, 2, 6])
+
+NORM_VEC_MASS = np.array(NORM_VEC.tolist() + [1471])
+
+# Assume 64 is the highest repeat of any 1 atom
+MAX_ELEMENT_NUM = 64
+
+element_to_ind = dict(zip(VALID_ELEMENTS, np.arange(len(VALID_ELEMENTS))))
+element_to_position = dict(zip(VALID_ELEMENTS, ELEMENT_VECTORS))
+element_to_position_mass = dict(zip(VALID_ELEMENTS, ELEMENT_VECTORS_MASS))
+
+ION_LST = [
+    "[M+H]+",
+    "[M+Na]+",
+    "[M+K]+",
+    "[M-H2O+H]+",
+    "[M+H3N+H]+",
+    "[M]+",
+    "[M-H4O2+H]+",
+    "[M-H]-"
+]
+
+ion_remap = dict(zip(ION_LST, ION_LST))
+ion_remap.update(
+    {
+        "[M+NH4]+": "[M+H3N+H]+",
+        "M+H": "[M+H]+",
+        "M+Na": "[M+Na]+",
+        "M+H-H2O": "[M-H2O+H]+",
+        "M-H2O+H": "[M-H2O+H]+",
+        "M+NH4": "[M+H3N+H]+",
+        "M-2H2O+H": "[M-H4O2+H]+",
+        "[M-2H2O+H]+": "[M-H4O2+H]+",
+        "[M-H]-": "[M-H]-",
+    }
+)
+
+ion_to_idx = dict(zip(ION_LST, np.arange(len(ION_LST))))
+
+ion_to_mass = {
+    "[M+H]+": ELEMENT_TO_MASS["H"] - ELECTRON_MASS,
+    "[M+Na]+": ELEMENT_TO_MASS["Na"] - ELECTRON_MASS,
+    "[M+K]+": ELEMENT_TO_MASS["K"] - ELECTRON_MASS,
+    "[M-H2O+H]+": -ELEMENT_TO_MASS["O"] - ELEMENT_TO_MASS["H"] - ELECTRON_MASS,
+    "[M+H3N+H]+": ELEMENT_TO_MASS["N"] + ELEMENT_TO_MASS["H"] * 4 - ELECTRON_MASS,
+    "[M]+": 0 - ELECTRON_MASS,
+    "[M-H4O2+H]+": -ELEMENT_TO_MASS["O"] * 2 - ELEMENT_TO_MASS["H"] * 3 - ELECTRON_MASS,
+    "[M-H]-": ELEMENT_TO_MASS["H"] + ELECTRON_MASS,
+}
+
+ion_to_add_vec = {
+    "[M+H]+": element_to_position["H"],
+    "[M+Na]+": element_to_position["Na"],
+    "[M+K]+": element_to_position["K"],
+    "[M-H2O+H]+": -element_to_position["O"] - element_to_position["H"],
+    "[M+H3N+H]+": element_to_position["N"] + element_to_position["H"] * 4,
+    "[M]+": np.zeros_like(element_to_position["H"]),
+    "[M-H4O2+H]+": -element_to_position["O"] * 2 - element_to_position["H"] * 3,
+}
+
+instrument_to_type = defaultdict(lambda : "unknown")
+instrument_to_type.update({
+    "Thermo Finnigan Velos Orbitrap": "orbitrap",
+    "Thermo Finnigan Elite Orbitrap": "orbitrap",
+    "Orbitrap Fusion Lumos": "orbitrap",
+    "Q-ToF (LCMS)": "qtof",
+    "Unknown (LCMS)": "unknown",
+    "ion trap": "iontrap",
+    "FTICR (LCMS)": "fticr",
+    "Bruker Q-ToF (LCMS)": "qtof",
+    "Orbitrap (LCMS)": "orbitrap",
+})
+
+instruments = sorted(list(set(instrument_to_type.values())))
+max_instr_idx = len(instruments) + 1
+instrument_to_idx = dict(zip(instruments, np.arange(len(instruments))))
+
+
+# Define rdbe mult
+rdbe_mult = np.zeros_like(ELEMENT_VECTORS[0])
+els = ["C", "N", "P", "H", "Cl", "Br", "I", "F"]
+weights = [2, 1, 1, -1, -1, -1, -1, -1]
+for k, v in zip(els, weights):
+    rdbe_mult[element_to_ind[k]] = v
+
+
+def get_ion_idx(ionization: str) -> int:
+    """map ionization to its index in one hot encoding"""
+    return ion_to_idx[ionization]
+
+
+def get_instr_idx(instrument: str) -> int:
+    """map instrument to its index in one hot encoding"""
+    inst = instrument_to_type.get(instrument, "unknown")
+    return instrument_to_idx[inst]
+
+
+def has_valid_els(chem_formula: str) -> bool:
+    """has_valid_els"""
+    for (chem_symbol, num) in re.findall(CHEM_FORMULA_SIZE, chem_formula):
+        if chem_symbol not in VALID_ELEMENTS:
+            return False
+    return True
+
+
+def formula_to_dense(chem_formula: str) -> np.ndarray:
+    """formula_to_dense.
+
+    Args:
+        chem_formula (str): Input chemical formal
+    Return:
+        np.ndarray of vector
+
+    """
+    total_onehot = []
+    for (chem_symbol, num) in re.findall(CHEM_FORMULA_SIZE, chem_formula):
+        # Convert num to int
+        num = 1 if num == "" else int(num)
+        one_hot = element_to_position[chem_symbol].reshape(1, -1)
+        one_hot_repeats = np.repeat(one_hot, repeats=num, axis=0)
+        total_onehot.append(one_hot_repeats)
+
+    # Check if null
+    if len(total_onehot) == 0:
+        dense_vec = np.zeros(len(element_to_position))
+    else:
+        dense_vec = np.vstack(total_onehot).sum(0)
+    return dense_vec
+
+
+def cross_sum(x, y):
+    """cross_sum."""
+    return (np.expand_dims(x, 0) + np.expand_dims(y, 1)).reshape(-1, y.shape[-1])
+
+
+def get_all_subsets_dense(
+    dense_formula: str, element_vectors
+) -> (np.ndarray, np.ndarray):
+    """_summary_
+
+    Args:
+        dense_formula (str, element_vectors): _description_
+        np (_type_): _description_
+
+    Returns:
+        _type_: _description_
+    """
+
+    non_zero = np.argwhere(dense_formula > 0).flatten()
+
+    vectorized_formula = []
+    for nonzero_ind in non_zero:
+        temp = element_vectors[nonzero_ind] * np.arange(
+            0, dense_formula[nonzero_ind] + 1
+        ).reshape(-1, 1)
+        vectorized_formula.append(temp)
+
+    zero_vec = np.zeros((1, element_vectors.shape[-1]))
+    cross_prod = reduce(cross_sum, vectorized_formula, zero_vec)
+
+    cross_prod_inds = rdbe_filter(cross_prod)
+    cross_prod = cross_prod[cross_prod_inds]
+    all_masses = cross_prod.dot(VALID_MONO_MASSES)
+    return cross_prod, all_masses
+
+
+def get_all_subsets(chem_formula: str):
+    dense_formula = formula_to_dense(chem_formula)
+    return get_all_subsets_dense(dense_formula, element_vectors=ELEMENT_VECTORS)
+
+
+def rdbe_filter(cross_prod):
+    """rdbe_filter.
+    Args:
+        cross_prod:
+    """
+    rdbe_total = 1 + 0.5 * cross_prod.dot(rdbe_mult)
+    filter_inds = np.argwhere(rdbe_total >= 0).flatten()
+    return filter_inds
+
+
+def formula_to_dense(chem_formula: str) -> np.ndarray:
+    """formula_to_dense.
+
+    Args:
+        chem_formula (str): Input chemical formal
+    Return:
+        np.ndarray of vector
+
+    """
+    total_onehot = []
+    for (chem_symbol, num) in re.findall(CHEM_FORMULA_SIZE, chem_formula):
+        # Convert num to int
+        num = 1 if num == "" else int(num)
+        one_hot = element_to_position[chem_symbol].reshape(1, -1)
+        one_hot_repeats = np.repeat(one_hot, repeats=num, axis=0)
+        total_onehot.append(one_hot_repeats)
+
+    # Check if null
+    if len(total_onehot) == 0:
+        dense_vec = np.zeros(len(element_to_position))
+    else:
+        dense_vec = np.vstack(total_onehot).sum(0)
+
+    return dense_vec
+
+
+def formula_to_dense_mass(chem_formula: str) -> np.ndarray:
+    """formula_to_dense_mass.
+
+    Return formula including full compound mass
+
+    Args:
+        chem_formula (str): Input chemical formal
+    Return:
+        np.ndarray of vector
+
+    """
+    total_onehot = []
+    for (chem_symbol, num) in re.findall(CHEM_FORMULA_SIZE, chem_formula):
+        # Convert num to int
+        num = 1 if num == "" else int(num)
+        one_hot = element_to_position_mass[chem_symbol].reshape(1, -1)
+        one_hot_repeats = np.repeat(one_hot, repeats=num, axis=0)
+        total_onehot.append(one_hot_repeats)
+
+    # Check if null
+    if len(total_onehot) == 0:
+        dense_vec = np.zeros(len(element_to_position_mass["H"]))
+    else:
+        dense_vec = np.vstack(total_onehot).sum(0)
+
+    return dense_vec
+
+
+def formula_to_dense_mass_norm(chem_formula: str) -> np.ndarray:
+    """formula_to_dense_mass_norm.
+
+    Return formula including full compound mass and normalized
+
+    Args:
+        chem_formula (str): Input chemical formal
+    Return:
+        np.ndarray of vector
+
+    """
+    dense_vec = formula_to_dense_mass(chem_formula)
+    dense_vec = dense_vec / NORM_VEC_MASS
+
+    return dense_vec
+
+
+def formula_mass(chem_formula: str) -> float:
+    """get formula mass"""
+    mass = 0
+    for (chem_symbol, num) in re.findall(CHEM_FORMULA_SIZE, chem_formula):
+        # Convert num to int
+        num = 1 if num == "" else int(num)
+        mass += ELEMENT_TO_MASS[chem_symbol] * num
+    return mass
+
+
+def electron_correct(mass: float) -> float:
+    """subtract the rest mass of an electron"""
+    return mass - ELECTRON_MASS
+
+
+def formula_difference(formula_1, formula_2):
+    """formula_1 - formula_2"""
+    form_1 = {
+        chem_symbol: (int(num) if num != "" else 1)
+        for chem_symbol, num in re.findall(CHEM_FORMULA_SIZE, formula_1)
+    }
+    form_2 = {
+        chem_symbol: (int(num) if num != "" else 1)
+        for chem_symbol, num in re.findall(CHEM_FORMULA_SIZE, formula_2)
+    }
+
+    for k, v in form_2.items():
+        form_1[k] = form_1[k] - form_2[k]
+    out_formula = "".join([f"{k}{v}" for k, v in form_1.items() if v > 0])
+    return out_formula
+
+
+def get_mol_from_structure_string(structure_string, structure_type):
+    if structure_type == "InChI":
+        mol = Chem.MolFromInchi(structure_string)
+    else:
+        mol = Chem.MolFromSmiles(structure_string)
+    return mol
+
+
+def vec_to_formula(form_vec):
+    """vec_to_formula."""
+    build_str = ""
+    for i in np.argwhere(form_vec > 0).flatten():
+        el = VALID_ELEMENTS[i]
+        ct = int(form_vec[i])
+        new_item = f"{el}{ct}" if ct > 1 else f"{el}"
+        build_str = build_str + new_item
+    return build_str
+
+
+def standardize_form(i):
+    """standardize_form."""
+    return vec_to_formula(formula_to_dense(i))
+
+
+def standardize_adduct(adduct):
+    """standardize_adduct."""
+    adduct = adduct.replace(" ", "")
+    adduct = ion_remap.get(adduct, adduct)
+    if adduct not in ION_LST:
+        raise ValueError(f"Adduct {adduct} not in ION_LST")
+    return adduct
+
+
+def calc_structure_string_type(structure_string):
+    """calc_structure_string_type.
+
+    Args:
+        structure_string:
+    """
+    structure_type = None
+    if pd.isna(structure_string):
+        structure_type = "empty"
+    elif structure_string.startswith("InChI="):
+        structure_type = "InChI"
+    elif Chem.MolFromSmiles(structure_string) is not None:
+        structure_type = "Smiles"
+    return structure_type
+
+
+def uncharged_formula(mol, mol_type="mol") -> str:
+    """Compute uncharged formula"""
+    if mol_type == "mol":
+        chem_formula = CalcMolFormula(mol)
+    elif mol_type == "smiles":
+        mol = Chem.MolFromSmiles(mol)
+        if mol is None:
+            return None
+        chem_formula = CalcMolFormula(mol)
+    else:
+        raise ValueError()
+
+    return re.findall(r"^([^\+,^\-]*)", chem_formula)[0]
+
+
+def form_from_smi(smi: str) -> str:
+    """form_from_smi.
+
+    Args:
+        smi (str): smi
+
+    Return:
+        str
+    """
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return ""
+    else:
+        return CalcMolFormula(mol)
+
+
+def inchikey_from_smiles(smi: str) -> str:
+    """inchikey_from_smiles.
+
+    Args:
+        smi (str): smi
+
+    Returns:
+        str:
+    """
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return ""
+    else:
+        return Chem.MolToInchiKey(mol)
+
+
+def contains_metals(formula: str) -> bool:
+    """returns true if formula contains metals"""
+    METAL_RE = "(Fe|Co|Zn|Rh|Pt|Li)"
+    return len(re.findall(METAL_RE, formula)) > 0
+
+
+class SmilesStandardizer(object):
+    """Standardize smiles"""
+
+    def __init__(self, *args, **kwargs):
+        self.fragment_standardizer = rdMolStandardize.LargestFragmentChooser()
+        self.charge_standardizer = rdMolStandardize.Uncharger()
+
+    def standardize_smiles(self, smi):
+        """Standardize smiles string"""
+        mol = Chem.MolFromSmiles(smi)
+        out_smi = self.standardize_mol(mol)
+        return out_smi
+
+    def standardize_mol(self, mol) -> str:
+        """Standardize smiles string"""
+        mol = self.fragment_standardizer.choose(mol)
+        mol = self.charge_standardizer.uncharge(mol)
+
+        # Round trip to and from inchi to tautomer correct
+        # Also standardize tautomer in the middle
+        output_smi = Chem.MolToSmiles(mol, isomericSmiles=False)
+        return output_smi
+
+
+def mass_from_smi(smi: str) -> float:
+    """mass_from_smi.
+
+    Args:
+        smi (str): smi
+
+    Return:
+        str
+    """
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return 0
+    else:
+        return ExactMolWt(mol)
+
+
+def min_formal_from_smi(smi: str):
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return 0
+    else:
+        formal = np.array([j.GetFormalCharge() for j in mol.GetAtoms()])
+        return formal.min()
+
+
+def max_formal_from_smi(smi: str):
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return 0
+    else:
+        formal = np.array([j.GetFormalCharge() for j in mol.GetAtoms()])
+        return formal.max()
+
+
+def atoms_from_smi(smi: str) -> int:
+    """atoms_from_smi.
+
+    Args:
+        smi (str): smi
+
+    Return:
+        int
+    """
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return 0
+    else:
+        return mol.GetNumAtoms()
+
+
+def has_valid_els(chem_formula: str) -> bool:
+    """has_valid_els"""
+    for (chem_symbol, num) in re.findall(CHEM_FORMULA_SIZE, chem_formula):
+        if chem_symbol not in VALID_ELEMENTS:
+            return False
+    return True
+
+
+def add_ion(form: str, ion: str):
+    """add_ion.
+    Args:
+        form (str): form
+        ion (str): ion
+    """
+    ion_vec = ion_to_add_vec[ion]
+    form_vec = formula_to_dense(form)
+    return vec_to_formula(form_vec + ion_vec)
+
+
+def achiral_smi(smi: str) -> str:
+    """achiral_smi.
+
+    Return:
+        isomeric smiles
+
+    """
+    try:
+        mol = Chem.MolFromSmiles(smi)
+        if mol is not None:
+            smi = Chem.MolToSmiles(mol, isomericSmiles=False)
+            return smi
+        else:
+            return ""
+    except:
+        return ""
+
+
+def npclassifer_query(inputs):
+    """npclassifier_query.
+
+    Args:
+        input: Tuple of name, molecule
+    Return:
+        Dict of name to molecule
+    """
+    import requests
+
+    spec = inputs[0]
+    endpoint = "https://npclassifier.ucsd.edu/classify"
+    req_data = {"smiles": inputs[1]}
+    out = requests.get(f"{endpoint}", data=req_data)
+    out.raise_for_status()
+    out_json = out.json()
+    return {spec: out_json}
+
+
+def clipped_ppm(mass_diff: np.ndarray, parentmass: np.ndarray) -> np.ndarray:
+    """clipped_ppm.
+
+    Args:
+        mass_diff (np.ndarray): mass_diff
+        parentmass (np.ndarray): parentmass
+
+    Returns:
+        np.ndarray:
+    """
+    parentmass_copy = parentmass * 1
+    parentmass_copy[parentmass < 200] = 200
+    ppm = mass_diff / parentmass_copy * 1e6
+    return ppm
+
+
+def clipped_ppm_single(
+    cls_mass_diff: float,
+    parentmass: float,
+):
+    """clipped_ppm_single.
+
+    Args:
+        cls_mass_diff (float): cls_mass_diff
+        parentmass (float): parentmass
+    """
+    div_factor = 200 if parentmass < 200 else parentmass
+    cls_ppm = cls_mass_diff / div_factor * 1e6
+    return cls_ppm

mvp/subformula_assign/utils/parallel_utils.py

@@ -0,0 +1,84 @@
+"""parallel_utils.py"""
+import logging
+from multiprocess.context import TimeoutError
+from pathos import multiprocessing as mp
+from tqdm import tqdm
+
+
+def simple_parallel(
+    input_list, function, max_cpu=16, timeout=4000, max_retries=3, use_ray: bool = False
+):
+    """Simple parallelization.
+
+    Use map async and retries in case we get odd stalling behavior.
+
+    input_list: Input list to op on
+    function: Fn to apply
+    max_cpu: Num cpus
+    timeout: Length of timeout
+    max_retries: Num times to retry this
+    use_ray
+
+    """
+    # If ray is required. Set to false.
+    if use_ray and False:
+        import ray
+
+        @ray.remote
+        def ray_func(x):
+            return function(x)
+
+        return ray.get([ray_func.remote(x) for x in input_list])
+
+    from multiprocess.context import TimeoutError
+    from pathos import multiprocessing as mp
+
+    cpus = min(mp.cpu_count(), max_cpu)
+    pool = mp.Pool(processes=cpus)
+    results = pool.map(function, input_list)
+    pool.close()
+    pool.join()
+    return results
+
+
+def chunked_parallel(
+    input_list, function, chunks=100, max_cpu=16, timeout=4000, max_retries=3
+):
+    """chunked_parallel.
+
+    Args:
+        input_list : list of objects to apply function
+        function : Callable with 1 input and returning a single value
+        chunks: number of hcunks
+        max_cpu: Max num cpus
+        timeout: Length of timeout
+        max_retries: Num times to retry this
+    """
+
+    # Adding it here fixes somessetting disrupted elsewhere
+
+    def batch_func(list_inputs):
+        outputs = []
+        for i in list_inputs:
+            outputs.append(function(i))
+        return outputs
+
+    list_len = len(input_list)
+    num_chunks = min(list_len, chunks)
+    step_size = len(input_list) // num_chunks + 1
+
+    chunked_list = [
+        input_list[i : i + step_size] for i in range(0, len(input_list), step_size)
+    ]
+
+    list_outputs = simple_parallel(
+        chunked_list,
+        batch_func,
+        max_cpu=max_cpu,
+        timeout=timeout,
+        max_retries=max_retries,
+    )
+    # Unroll
+    full_output = [j for i in list_outputs for j in i]
+
+    return full_output

mvp/subformula_assign/utils/parse_utils.py

@@ -0,0 +1,295 @@
+""" parse_utils.py """
+from pathlib import Path
+from typing import Tuple, List, Optional
+from itertools import groupby
+
+from tqdm import tqdm
+import numpy as np
+import pandas as pd
+
+
+def parse_spectra(spectra_file: str) -> Tuple[dict, List[Tuple[str, np.ndarray]]]:
+    """parse_spectra.
+
+    Parses spectra in the SIRIUS format and returns
+
+    Args:
+        spectra_file (str): Name of spectra file to parse
+    Return:
+        Tuple[dict, List[Tuple[str, np.ndarray]]]: metadata and list of spectra
+            tuples containing name and array
+    """
+    lines = [i.strip() for i in open(spectra_file, "r").readlines()]
+
+    group_num = 0
+    metadata = {}
+    spectras = []
+    my_iterator = groupby(
+        lines, lambda line: line.startswith(">") or line.startswith("#")
+    )
+
+    for index, (start_line, lines) in enumerate(my_iterator):
+        group_lines = list(lines)
+        subject_lines = list(next(my_iterator)[1])
+        # Get spectra
+        if group_num > 0:
+            spectra_header = group_lines[0].split(">")[1]
+            peak_data = [
+                [float(x) for x in peak.split()[:2]]
+                for peak in subject_lines
+                if peak.strip()
+            ]
+            # Check if spectra is empty
+            if len(peak_data):
+                peak_data = np.vstack(peak_data)
+                # Add new tuple
+                spectras.append((spectra_header, peak_data))
+        # Get meta data
+        else:
+            entries = {}
+            for i in group_lines:
+                if " " not in i:
+                    continue
+                elif i.startswith("#INSTRUMENT TYPE"):
+                    key = "#INSTRUMENT TYPE"
+                    val = i.split(key)[1].strip()
+                    entries[key[1:]] = val
+                else:
+                    start, end = i.split(" ", 1)
+                    start = start[1:]
+                    while start in entries:
+                        start = f"{start}'"
+                    entries[start] = end
+
+            metadata.update(entries)
+        group_num += 1
+
+    metadata["_FILE_PATH"] = spectra_file
+    metadata["_FILE"] = Path(spectra_file).stem
+    return metadata, spectras
+
+
+def spec_to_ms_str(
+    spec: List[Tuple[str, np.ndarray]], essential_keys: dict, comments: dict = {}
+) -> str:
+    """spec_to_ms_str.
+
+    Turn spec ars and info dicts into str for output file
+
+
+    Args:
+        spec (List[Tuple[str, np.ndarray]]): spec
+        essential_keys (dict): essential_keys
+        comments (dict): comments
+
+    Returns:
+        str:
+    """
+
+    def pair_rows(rows):
+        return "\n".join([f"{i} {j}" for i, j in rows])
+
+    header = "\n".join(f">{k} {v}" for k, v in essential_keys.items())
+    comments = "\n".join(f"#{k} {v}" for k, v in essential_keys.items())
+    spec_strs = [f">{name}\n{pair_rows(ar)}" for name, ar in spec]
+    spec_str = "\n\n".join(spec_strs)
+    output = f"{header}\n{comments}\n\n{spec_str}"
+    return output
+
+
+def build_mgf_str(
+    meta_spec_list: List[Tuple[dict, List[Tuple[str, np.ndarray]]]],
+    merge_charges=True,
+    parent_mass_keys=["PEPMASS", "parentmass", "PRECURSOR_MZ"],
+) -> str:
+    """build_mgf_str.
+
+    Args:
+        meta_spec_list (List[Tuple[dict, List[Tuple[str, np.ndarray]]]]): meta_spec_list
+
+    Returns:
+        str:
+    """
+    entries = []
+    for meta, spec in tqdm(meta_spec_list):
+        str_rows = ["BEGIN IONS"]
+
+        # Try to add precusor mass
+        for i in parent_mass_keys:
+            if i in meta:
+                pep_mass = float(meta.get(i, -100))
+                str_rows.append(f"PEPMASS={pep_mass}")
+                break
+
+        for k, v in meta.items():
+            str_rows.append(f"{k.upper().replace(' ', '_')}={v}")
+
+        if merge_charges:
+            spec_ar = np.vstack([i[1] for i in spec])
+            spec_ar = np.vstack([i for i in sorted(spec_ar, key=lambda x: x[0])])
+        else:
+            raise NotImplementedError()
+        str_rows.extend([f"{i} {j}" for i, j in spec_ar])
+        str_rows.append("END IONS")
+
+        str_out = "\n".join(str_rows)
+        entries.append(str_out)
+
+    full_out = "\n\n".join(entries)
+    return full_out
+
+
+def parse_spectra_msp(
+    mgf_file: str, max_num: Optional[int] = None
+) -> List[Tuple[dict, List[Tuple[str, np.ndarray]]]]:
+    """parse_spectr_msp.
+
+    Parses spectra in the MSP file format
+
+    Args:
+        mgf_file (str) : str
+        max_num (Optional[int]): If set, only parse this many
+    Return:
+        List[Tuple[dict, List[Tuple[str, np.ndarray]]]]: metadata and list of spectra
+            tuples containing name and array
+    """
+
+    key = lambda x: x.strip().startswith("PEPMASS")
+    parsed_spectra = []
+    with open(mgf_file, "r", encoding="utf-8") as fp:
+        for (is_header, group) in tqdm(groupby(fp, key)):
+
+            if is_header:
+                continue
+            meta = dict()
+            spectra = []
+            # Note: Sometimes we have multiple scans
+            # This mgf has them collapsed
+            cur_spectra_name = "spec"
+            cur_spectra = []
+            group = list(group)
+            for line in group:
+                line = line.strip()
+                if not line:
+                    pass
+                elif ":" in line:
+                    k, v = [i.strip() for i in line.split(":", 1)]
+                    meta[k] = v
+                else:
+                    mz, intens = line.split()
+                    cur_spectra.append((float(mz), float(intens)))
+
+            if len(cur_spectra) > 0:
+                cur_spectra = np.vstack(cur_spectra)
+                spectra.append((cur_spectra_name, cur_spectra))
+                parsed_spectra.append((meta, spectra))
+            else:
+                pass
+                # print("no spectra found for group: ", "".join(group))
+
+            if max_num is not None and len(parsed_spectra) > max_num:
+                # print("Breaking")
+                break
+        return parsed_spectra
+
+
+def parse_spectra_mgf(
+    mgf_file: str, max_num: Optional[int] = None
+) -> List[Tuple[dict, List[Tuple[str, np.ndarray]]]]:
+    """parse_spectr_mgf.
+
+    Parses spectra in the MGF file formate, with
+
+    Args:
+        mgf_file (str) : str
+        max_num (Optional[int]): If set, only parse this many
+    Return:
+        List[Tuple[dict, List[Tuple[str, np.ndarray]]]]: metadata and list of spectra
+            tuples containing name and array
+    """
+
+    key = lambda x: x.strip() == "BEGIN IONS"
+    parsed_spectra = []
+    with open(mgf_file, "r") as fp:
+
+        for (is_header, group) in tqdm(groupby(fp, key)):
+
+            if is_header:
+                continue
+
+            meta = dict()
+            spectra = []
+            # Note: Sometimes we have multiple scans
+            # This mgf has them collapsed
+            cur_spectra_name = "spec"
+            cur_spectra = []
+            group = list(group)
+            for line in group:
+                line = line.strip()
+                if not line:
+                    pass
+                elif line == "END IONS" or line == "BEGIN IONS":
+                    pass
+                elif "=" in line:
+                    k, v = [i.strip() for i in line.split("=", 1)]
+                    meta[k] = v
+                else:
+                    mz, intens = line.split()
+                    cur_spectra.append((float(mz), float(intens)))
+
+            if len(cur_spectra) > 0:
+                cur_spectra = np.vstack(cur_spectra)
+                spectra.append((cur_spectra_name, cur_spectra))
+                parsed_spectra.append((meta, spectra))
+            else:
+                pass
+                # print("no spectra found for group: ", "".join(group))
+
+            if max_num is not None and len(parsed_spectra) > max_num:
+                # print("Breaking")
+                break
+        return parsed_spectra
+
+
+def parse_tsv_spectra(spectra_file: str) -> List[Tuple[str, np.ndarray]]:
+    """parse_tsv_spectra.
+
+    Parses spectra returned from sirius fragmentation tree
+
+    Args:
+        spectra_file (str): Name of spectra tsv file to parse
+    Return:
+        List[Tuple[str, np.ndarray]]]: list of spectra
+            tuples containing name and array. This is used to maintain
+            consistency with the parse_spectra output
+    """
+    output_spec = []
+    with open(spectra_file, "r") as fp:
+        for index, line in enumerate(fp):
+            if index == 0:
+                continue
+            line = line.strip().split("\t")
+            intensity = float(line[1])
+            exact_mass = float(line[3])
+            output_spec.append([exact_mass, intensity])
+
+    output_spec = np.array(output_spec)
+    return_obj = [("sirius_spec", output_spec)]
+    return return_obj
+
+# YZC parse msgym-like formatted data
+def parse_spectra_msgym(df):
+    
+    parsed_spectra = []
+    for _, row in df.iterrows():
+        mzs = [float(m) for m in row['mzs'].split(',')]
+        intensities = [float(i) for i in row['intensities'].split(',')]
+        cur_spectra = [(m, i) for m, i in zip(mzs, intensities)]
+        cur_spectra = np.vstack(cur_spectra)
+        cur_spectra_name = row['spec']
+        meta = {'ID': cur_spectra_name,
+        'parentmass': row['parent_mass']}
+        parsed_spectra.append((meta, [(cur_spectra_name, cur_spectra)]))
+    return parsed_spectra
+
+

mvp/subformula_assign/utils/spectra_utils.py

@@ -0,0 +1,325 @@
+""" spectra_utils.py"""
+import logging
+import numpy as np
+from typing import List
+
+
+from .chem_utils import (
+    vec_to_formula,
+    get_all_subsets,
+    ion_to_mass,
+    ION_LST,
+    clipped_ppm,
+)
+
+
+def bin_spectra(
+    spectras: List[np.ndarray], num_bins: int = 2000, upper_limit: int = 1000
+) -> np.ndarray:
+    """bin_spectra.
+
+    Args:
+        spectras (List[np.ndarray]): Input list of spectra tuples
+            [(header, spec array)]
+        num_bins (int): Number of discrete bins from [0, upper_limit)
+        upper_limit (int): Max m/z to consider featurizing
+
+    Return:
+        np.ndarray of shape [channels, num_bins]
+    """
+    bins = np.linspace(0, upper_limit, num=num_bins)
+    binned_spec = np.zeros((len(spectras), len(bins)))
+    for spec_index, spec in enumerate(spectras):
+
+        # Convert to digitized spectra
+        digitized_mz = np.digitize(spec[:, 0], bins=bins)
+
+        # Remove all spectral peaks out of range
+        in_range = digitized_mz < len(bins)
+        digitized_mz, spec = digitized_mz[in_range], spec[in_range, :]
+
+        # Add the current peaks to the spectra
+        # Use a loop rather than vectorize because certain bins have conflicts
+        # based upon resolution
+        for bin_index, spec_val in zip(digitized_mz, spec[:, 1]):
+            binned_spec[spec_index, bin_index] += spec_val
+
+    return binned_spec
+
+
+def merge_norm_spectra(spec_tuples, precision=4) -> np.ndarray:
+    """merge_norm_spectra.
+
+    Take a list of mz, inten tuple arrays and merge them by 4 digit precision
+
+    Note this uses _max_ merging
+
+    """
+    mz_to_inten_pair = {}
+    for i in spec_tuples:
+        for tup in i:
+            mz, inten = tup
+            mz_ind = np.round(mz, precision)
+            cur_pair = mz_to_inten_pair.get(mz_ind)
+            if cur_pair is None:
+                mz_to_inten_pair[mz_ind] = tup
+            elif inten > cur_pair[1]:
+                mz_to_inten_pair[mz_ind] = (mz_ind, inten)
+            else:
+                pass
+
+    merged_spec = np.vstack([v for k, v in mz_to_inten_pair.items()])
+    merged_spec[:, 1] = merged_spec[:, 1] / merged_spec[:, 1].max()
+    return merged_spec
+
+
+def norm_spectrum(binned_spec: np.ndarray) -> np.ndarray:
+    """norm_spectrum.
+
+    Normalizes each spectral channel to have norm 1
+    This change is made in place
+
+    Args:
+        binned_spec (np.ndarray) : Vector of spectras
+
+    Return:
+        np.ndarray where each channel has max(1)
+    """
+
+    spec_maxes = binned_spec.max(1)
+
+    non_zero_max = spec_maxes > 0
+
+    spec_maxes = spec_maxes[non_zero_max]
+    binned_spec[non_zero_max] = binned_spec[non_zero_max] / spec_maxes.reshape(-1, 1)
+
+    return binned_spec
+
+
+def process_spec_file(meta, tuples, precision=4, max_inten=0.001, max_peaks=60):
+    """process_spec_file."""
+
+    if "parentmass" in meta:
+        parentmass = meta.get("parentmass", None)
+    elif "PARENTMASS" in meta:
+        parentmass = meta.get("PARENTMASS", None)
+    elif "PEPMASS" in meta:
+        parentmass = meta.get("PEPMASS", None)
+    else:
+        logging.debug(f"missing parentmass for spec")
+        parentmass = 1000000
+
+    parentmass = float(parentmass)
+
+    # First norm spectra
+    fused_tuples = [x for _, x in tuples if x.size > 0]
+
+    if len(fused_tuples) == 0:
+        return
+
+    mz_to_inten_pair = {}
+    new_tuples = []
+    for i in fused_tuples:
+        for tup in i:
+            mz, inten = tup
+            mz_ind = np.round(mz, precision)
+            cur_pair = mz_to_inten_pair.get(mz_ind)
+            if cur_pair is None:
+                mz_to_inten_pair[mz_ind] = tup
+                new_tuples.append(tup)
+            elif inten > cur_pair[1]:
+                cur_pair[1] = inten
+            else:
+                pass
+
+    merged_spec = np.vstack(new_tuples)
+    merged_spec = merged_spec[merged_spec[:, 0] <= (parentmass + 1)] # could end up removing all peaks?
+    try:
+        merged_spec[:, 1] = merged_spec[:, 1] / merged_spec[:, 1].max()
+    except:
+        return
+
+    # Sqrt intensities here
+    merged_spec[:, 1] = np.sqrt(merged_spec[:, 1])
+
+    merged_spec = max_inten_spec(
+        merged_spec, max_num_inten=max_peaks, inten_thresh=max_inten
+    )
+    return merged_spec
+
+
+def max_inten_spec(spec, max_num_inten: int = 60, inten_thresh: float = 0):
+    """max_inten_spec.
+
+    Args:
+        spec: 2D spectra array
+        max_num_inten: Max number of peaks
+        inten_thresh: Min intensity to alloow in returned peak
+
+    Return:
+        Spec filtered down
+
+
+    """
+    spec_masses, spec_intens = spec[:, 0], spec[:, 1]
+
+    # Make sure to only take max of each formula
+    # Sort by intensity and select top subpeaks
+    new_sort_order = np.argsort(spec_intens)[::-1]
+    if max_num_inten is not None:
+        new_sort_order = new_sort_order[:max_num_inten]
+
+    spec_masses = spec_masses[new_sort_order]
+    spec_intens = spec_intens[new_sort_order]
+
+    spec_mask = spec_intens > inten_thresh
+    spec_masses = spec_masses[spec_mask]
+    spec_intens = spec_intens[spec_mask]
+    spec = np.vstack([spec_masses, spec_intens]).transpose(1, 0)
+    return spec
+
+
+def max_thresh_spec(spec: np.ndarray, max_peaks=100, inten_thresh=0.003):
+    """max_thresh_spec.
+
+    Args:
+        spec (np.ndarray): spec
+        max_peaks: Max num peaks to keep
+        inten_thresh: Min inten to keep
+    """
+
+    spec_masses, spec_intens = spec[:, 0], spec[:, 1]
+
+    # Make sure to only take max of each formula
+    # Sort by intensity and select top subpeaks
+    new_sort_order = np.argsort(spec_intens)[::-1]
+    new_sort_order = new_sort_order[:max_peaks]
+
+    spec_masses = spec_masses[new_sort_order]
+    spec_intens = spec_intens[new_sort_order]
+
+    spec_mask = spec_intens > inten_thresh
+    spec_masses = spec_masses[spec_mask]
+    spec_intens = spec_intens[spec_mask]
+    out_ar = np.vstack([spec_masses, spec_intens]).transpose(1, 0)
+    return out_ar
+
+
+def assign_subforms(form, spec, ion_type, mass_diff_thresh=15):
+    """_summary_
+
+    Args:
+        form (_type_): _description_
+        spec (_type_): _description_
+        ion_type (_type_): _description_
+        mass_diff_thresh (int, optional): _description_. Defaults to 15.
+
+    Returns:
+        _type_: _description_
+    """
+    try:
+        cross_prod, masses = get_all_subsets(form)
+        spec_masses, spec_intens = spec[:, 0], spec[:, 1]
+
+        ion_masses = ion_to_mass[ion_type]
+        masses_with_ion = masses + ion_masses
+        ion_types = np.array([ion_type] * len(masses_with_ion))
+
+        mass_diffs = np.abs(spec_masses[:, None] - masses_with_ion[None, :])
+
+        formula_inds = mass_diffs.argmin(-1)
+        min_mass_diff = mass_diffs[np.arange(len(mass_diffs)), formula_inds]
+        rel_mass_diff = clipped_ppm(min_mass_diff, spec_masses)
+
+        # Filter by mass diff threshold (ppm)
+        valid_mask = rel_mass_diff < mass_diff_thresh
+        spec_masses = spec_masses[valid_mask]
+        spec_intens = spec_intens[valid_mask]
+        min_mass_diff = min_mass_diff[valid_mask]
+        rel_mass_diff = rel_mass_diff[valid_mask]
+        formula_inds = formula_inds[valid_mask]
+
+        formulas = np.array([vec_to_formula(j) for j in cross_prod[formula_inds]])
+        formula_masses = masses_with_ion[formula_inds]
+        ion_types = ion_types[formula_inds]
+
+        # Build mask for uniqueness on formula and ionization
+        # note that ionization are all the same for one subformula assignment
+        # hence we only need to consider the uniqueness of the formula
+        formula_idx_dict = {}
+        uniq_mask = []
+        for idx, formula in enumerate(formulas):
+            uniq_mask.append(formula not in formula_idx_dict)
+            gather_ind = formula_idx_dict.get(formula, None)
+            if gather_ind is None:
+                continue
+            spec_intens[gather_ind] += spec_intens[idx]
+            formula_idx_dict[formula] = idx
+
+        spec_masses = spec_masses[uniq_mask]
+        spec_intens = spec_intens[uniq_mask]
+        min_mass_diff = min_mass_diff[uniq_mask]
+        rel_mass_diff = rel_mass_diff[uniq_mask]
+        formula_masses = formula_masses[uniq_mask]
+        formulas = formulas[uniq_mask]
+        ion_types = ion_types[uniq_mask]
+
+        # To calculate explained intensity, preserve the original normalized
+        # intensity
+        if spec_intens.size == 0:
+            output_tbl = None
+        else:
+            output_tbl = {
+                "mz": list(spec_masses),
+                "ms2_inten": list(spec_intens),
+                "mono_mass": list(formula_masses),
+                "abs_mass_diff": list(min_mass_diff),
+                "mass_diff": list(rel_mass_diff),
+                "formula": list(formulas),
+                "ions": list(ion_types),
+            }
+    except:
+        output_tbl = None
+        print(f"failed to process formula {form}")
+        pass
+    output_dict = {
+        "cand_form": form,
+        "cand_ion": ion_type,
+        "output_tbl": output_tbl,
+    }
+    return output_dict
+
+
+def get_output_dict(
+    spec_name: str,
+    spec: np.ndarray,
+    form: str,
+    mass_diff_type: str,
+    mass_diff_thresh: float,
+    ion_type: str,
+) -> dict:
+    """_summary_
+
+    This function attemps to take an array of mass intensity values and assign
+    formula subsets to subpeaks
+
+    Args:
+        spec_name (str): _description_
+        spec (np.ndarray): _description_
+        form (str): _description_
+        mass_diff_type (str): _description_
+        mass_diff_thresh (float): _description_
+        ion_type (str): _description_
+
+    Returns:
+        dict: _description_
+    """
+    assert mass_diff_type == "ppm"
+    # This is the case for some erroneous MS2 files for which proc_spec_file return None
+    # All the MS2 subpeaks in these erroneous MS2 files has mz larger than parentmass
+    output_dict = {"cand_form": form, "cand_ion": ion_type, "output_tbl": None}
+    if spec is not None and ion_type in ION_LST:
+        output_dict = assign_subforms(
+            form, spec, ion_type, mass_diff_thresh=mass_diff_thresh
+        )
+    return output_dict

mvp/test.py

@@ -0,0 +1,118 @@
+import argparse
+import datetime
+import sys
+sys.path.insert(0, "/data/yzhouc01/MassSpecGym")
+sys.path.insert(0, "/data/yzhouc01/MVP")
+
+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,226 @@
+import os
+import json
+import numpy as np
+
+from mvp.data.transforms import SpecBinner, SpecBinnerLog, SpecFormulaFeaturizer, SpecFormulaMzFeaturizer, SpecMzIntTokenizer
+from massspecgym.data.transforms import SpecTransform, MolTransform
+from mvp.data.transforms import MolToGraph
+import mvp.data.datasets as jestr_datasets
+import typing as T
+from mvp.definitions import MSGYM_FORMULA_VECTOR_NORM, MSGYM_STANDARD_MH
+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,
+        "SpecFormulaMz": SpecFormulaMzFeaturizer,
+        'SpecMzIntTokens': SpecMzIntTokenizer}
+
+    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", "SpecFormulaMz"]:
+            featurizer_params.update({'element_list': params['element_list'], 'add_intensities': params['add_intensities'], 'formula_normalize_vector': MSGYM_FORMULA_VECTOR_NORM})
+
+        if view in ("SpecFormulaMz", 'SpecMzIntTokens'):
+            featurizer_params.update({'mz_mean_std': MSGYM_STANDARD_MH, 'mask_precursor': params['mask_precursor']})
+            # featurizer_params.update({'mask_precursor': params['mask_precursor']}) 
+        
+        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):
+    
+    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 'use_NL_spec' in params and params['use_NL_spec']:
+        dataset_params.update({'NL_spec_dir_pth': params['NL_spec_dir_pth']})
+    if params['pred_fp'] or params['use_fp']:
+        dataset_params.update({'fp_dir_pth': '', 'fp_size': params['fp_size'], 'fp_radius': params['fp_radius']})
+
+    return jestr_datasets.ExpandedRetrievalDataset(use_formulas=use_formulas, **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['aug_cands']:
+        dataset_params.update({'aug_cands_dir_pth': params['aug_cands_dir_pth'],
+                               'use_formulas':use_formulas,
+                               "aug_cands_size": params['aug_cands_size']})
+        
+    if params['use_cons_spec']:
+        dataset_params.update({'cons_spec_dir_pth': params['cons_spec_dir_pth']})
+
+    if 'use_NL_spec' in params and params['use_NL_spec']:
+        dataset_params.update({'NL_spec_dir_pth': params['NL_spec_dir_pth']})
+
+    # 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/debug.py

@@ -0,0 +1,19 @@
+import torch
+
+def nan_hook(self,inp, output):
+
+    nan_mask = torch.isnan(output)
+
+    if nan_mask.any():
+
+        print("In", self.__class__.__name__)
+
+        raise RuntimeError(f"Found NAN in output at indices: ", nan_mask.nonzero())
+
+    inf_mask = torch.isinf(output)
+
+    if inf_mask.any():
+
+        print("In", self.__class__.__name__)
+
+        raise RuntimeError(f"Found INF in output at indices: ", inf_mask.nonzero())

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,87 @@
+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
+
+
+def filip_similarity_batch(image_tokens, text_tokens, mask_image, mask_text):
+    """
+    Compute FILIP similarity for batches of image and text token embeddings.
+
+    Args:
+        image_tokens: (B, N_img, D) float tensor
+        text_tokens:  (B, N_text, D) float tensor
+        mask_image:   (B, N_img) bool tensor
+        mask_text:    (B, N_text) bool tensor
+
+    Returns:
+        similarities: (B,) float tensor of similarity scores
+    """
+    B, N_img, D = image_tokens.shape
+    N_text = text_tokens.shape[1]
+
+    # Normalize tokens
+    image_norm = F.normalize(image_tokens, p=2, dim=-1)  # (B, N_img, D)
+    text_norm = F.normalize(text_tokens, p=2, dim=-1)    # (B, N_text, D)
+
+    # Compute batched cosine similarity matrices
+    # Result shape: (B, N_img, N_text)
+    sim_matrix = torch.bmm(image_norm, text_norm.transpose(1, 2))
+
+    # Expand masks for broadcasting
+    mask_image_exp = mask_image.unsqueeze(2)  # (B, N_img, 1)
+    mask_text_exp = mask_text.unsqueeze(1)    # (B, 1, N_text)
+    valid_mask = mask_image_exp & mask_text_exp  # (B, N_img, N_text)
+
+    # Mask invalid positions by setting them to -inf
+    sim_matrix_masked = sim_matrix.masked_fill(~valid_mask, float('-inf'))
+
+    # Max over text tokens per image token: (B, N_img)
+    max_sim_img, _ = sim_matrix_masked.max(dim=2)
+
+    # Max over image tokens per text token: (B, N_text)
+    max_sim_text, _ = sim_matrix_masked.max(dim=1)
+
+    # Replace -inf (no valid tokens) with zeros to avoid NaNs
+    max_sim_img[max_sim_img == float('-inf')] = 0
+    max_sim_text[max_sim_text == float('-inf')] = 0
+
+    # Sum over valid tokens and divide by number of valid tokens (avoid division by zero)
+    sum_img = (max_sim_img * mask_image).sum(dim=1)
+    count_img = mask_image.sum(dim=1).clamp(min=1).float()
+
+    sum_text = (max_sim_text * mask_text).sum(dim=1)
+    count_text = mask_text.sum(dim=1).clamp(min=1).float()
+
+    avg_img = sum_img / count_img
+    avg_text = sum_text / count_text
+
+    # Final similarity per batch element
+    similarity = (avg_img + avg_text) / 2
+
+    return similarity

mvp/utils/loss.py

@@ -0,0 +1,156 @@
+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)
+
+
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+
+# ---------- Utility ----------
+def _safe_divide(num, denom, eps=1e-8):
+    return num / (denom + eps)
+
+
+# ---------- Single-GPU masked FILIP ----------
+def filip_loss_with_mask(a_tokens, b_tokens, mask_a, mask_b, temperature=0.07):
+    """
+    Single-GPU FILIP loss for modality A (spectra peaks) and modality B (graph nodes),
+    accounting for padding masks.
+
+    Args:
+        a_tokens: (B, N_a, D)   float tensor (will be normalized to unit vectors)
+        b_tokens: (B, N_b, D)
+        mask_a:   (B, N_a)      bool or byte tensor (True=valid)
+        mask_b:   (B, N_b)      bool or byte tensor
+        temperature: scalar or 0-dim tensor (learnable ok)
+
+    Returns:
+        scalar loss
+    """
+    device = a_tokens.device
+    B, N_a, D = a_tokens.shape
+    N_b = b_tokens.shape[1]
+
+    # normalize to cos sim
+    a = F.normalize(a_tokens, dim=-1)
+    b = F.normalize(b_tokens, dim=-1)
+
+    # Expand to compute all pairwise (batch-wise) similarities:
+    # sim shape: (B, B, N_a, N_b)  where sim[i,j,k,l] = dot(a[i,k], b[j,l])
+    a_exp = a.unsqueeze(1).expand(-1, B, -1, -1)  # (B, B, N_a, D)
+    b_exp = b.unsqueeze(0).expand(B, -1, -1, -1)  # (B, B, N_b, D)
+    sim = torch.einsum('bijd,bitd->bijt', a_exp, b_exp)  # (B, B, N_a, N_b)
+
+    # Expand masks to (B, B, N_a) and (B, B, N_b)
+    mask_a_exp = mask_a.unsqueeze(1).expand(-1, B, -1)  # (B, B, N_a)
+    mask_b_exp = mask_b.unsqueeze(0).expand(B, -1, -1)  # (B, B, N_b)
+
+    # ---- A -> B similarity (s_a2b) ----
+    # For every a-token we need max over valid b-tokens.
+    # Set invalid positions in sim to -inf before max.
+    sim_a2b = sim.clone()
+    invalid_b = ~mask_b_exp.unsqueeze(2).expand(-1, -1, sim_a2b.size(2), -1)  # (B, B, N_a, N_b)
+    sim_a2b[invalid_b] = float('-inf')
+
+    # max over b tokens -> (B, B, N_a)
+    max_over_b = sim_a2b.max(dim=3).values
+
+    # zero-out padded a-tokens then average over valid tokens
+    max_over_b = max_over_b * mask_a_exp  # padded a tokens get zero
+    denom_a = mask_a_exp.sum(dim=2).clamp(min=1).to(sim.dtype)  # (B, B)
+    s_a2b = max_over_b.sum(dim=2) / denom_a  # (B, B)
+
+    # ---- B -> A similarity (s_b2a) ----
+    sim_b2a = sim.clone()
+    invalid_a = ~mask_a_exp.unsqueeze(3).expand(-1,-1,-1,sim_b2a.size(3))  # (B, B, N_a, N_b)
+    sim_b2a[invalid_a] = float('-inf')
+
+    max_over_a = sim_b2a.max(dim=2).values  # (B, B, N_b)
+    max_over_a = max_over_a * mask_b_exp
+    denom_b = mask_b_exp.sum(dim=2).clamp(min=1).to(sim.dtype)
+    s_b2a = max_over_a.sum(dim=2) / denom_b  # (B, B)
+
+    # logits and loss
+    logits_a2b = s_a2b / temperature
+    logits_b2a = s_b2a / temperature
+
+    labels = torch.arange(B, device=device, dtype=torch.long)
+    loss_a2b = F.cross_entropy(logits_a2b, labels)
+    loss_b2a = F.cross_entropy(logits_b2a, labels)
+
+    return 0.5 * (loss_a2b + loss_b2a)
+

mvp/utils/models.py

@@ -0,0 +1,51 @@
+from mvp.models.spec_encoder import SpecEncMLP_BIN, SpecFormulaEncMLP, SpecFormulaTransformer,SpecFormula_mz_Encoder, SpecMzIntTokenTransformer
+from mvp.models.mol_encoder import MolEnc
+from mvp.models.encoders import MLP
+from mvp.models.contrastive import ContrastiveModel, CrossAttenContrastive, IndSpecEncoder, MultiViewContrastive, MultiViewFineTuning, FilipContrastive
+
+def get_spec_encoder(spec_enc:str, args):
+    return {"MLP_BIN": SpecEncMLP_BIN,
+            "MLP_Formula":SpecFormulaEncMLP,
+            "Transformer_Formula": SpecFormulaTransformer,
+            "Formula_BinnedSpec": SpecFormula_mz_Encoder,
+            "Transformer_MzInt": SpecMzIntTokenTransformer}[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 =='crossAttenContrastive':
+        model = CrossAttenContrastive(**params)
+    elif model == 'IndSpecEncoder':
+        params['pred_fp'] = False
+        params['use_cons_spec'] = False
+        model = IndSpecEncoder(**params)
+    elif model == "MultiviewContrastive":
+        model = MultiViewContrastive(**params)
+    elif model == "MultiViewFineTuning":
+        model = MultiViewFineTuning(**params)
+    elif model == "filipContrastive":
+        model = FilipContrastive(**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