msfiddle 0.1.0__py3-none-any.whl

msfiddle/download.py

@@ -0,0 +1,194 @@
+"""
+Module for downloading and locating pre-trained models for msfiddle.
+"""
+
+import os
+import sys
+import shutil
+import tempfile
+import urllib.request
+import zipfile
+import argparse
+from pathlib import Path
+from tqdm import tqdm
+
+
+MODEL_URLS = {
+    "fiddle_tcn_qtof": "https://github.com/JosieHong/FIDDLE/releases/download/v1.0.0/fiddle_tcn_qtof.zip",
+    "fiddle_fdr_qtof": "https://github.com/JosieHong/FIDDLE/releases/download/v1.0.0/fiddle_fdr_qtof.zip",
+    "fiddle_tcn_orbitrap": "https://github.com/JosieHong/FIDDLE/releases/download/v1.0.0/fiddle_tcn_orbitrap.zip",
+    "fiddle_fdr_orbitrap": "https://github.com/JosieHong/FIDDLE/releases/download/v1.0.0/fiddle_fdr_orbitrap.zip"
+}
+
+
+class DownloadProgressBar(tqdm):
+    """Progress bar for downloads using tqdm."""
+    
+    def update_to(self, b=1, bsize=1, tsize=None):
+        if tsize is not None:
+            self.total = tsize
+        self.update(b * bsize - self.n)
+
+
+def download_url(url, output_path):
+    """Download a file with a progress bar."""
+    with DownloadProgressBar(unit='B', unit_scale=True, miniters=1, desc=url.split('/')[-1]) as t:
+        urllib.request.urlretrieve(url, filename=output_path, reporthook=t.update_to)
+
+
+def get_package_dir():
+    """Get the directory where the package is installed."""
+    import msfiddle
+    return os.path.dirname(os.path.abspath(msfiddle.__file__))
+
+
+def get_checkpoint_dir():
+    """Get the directory where model checkpoints are stored."""
+    package_dir = get_package_dir()
+    return os.path.join(package_dir, "check_point")
+
+
+def download_models(destination=None, models=None):
+    """
+    Download pre-trained models for msfiddle.
+    
+    Args:
+        destination (str, optional): Path to save the models. If None, uses the package directory.
+        models (list, optional): List of model names to download. If None, downloads all models.
+    """
+    if destination is None:
+        destination = get_checkpoint_dir()
+    else:
+        destination = os.path.abspath(destination)
+    
+    # Create the destination directory if it doesn't exist
+    os.makedirs(destination, exist_ok=True)
+    
+    print(f"Models will be downloaded to: {destination}")
+    
+    # Determine which models to download
+    if models is None:
+        models = list(MODEL_URLS.keys())
+    
+    # Create a temporary directory for downloads
+    with tempfile.TemporaryDirectory() as temp_dir:
+        for model_name in models:
+            if model_name not in MODEL_URLS:
+                print(f"Warning: Unknown model '{model_name}', skipping.")
+                continue
+            
+            url = MODEL_URLS[model_name]
+            zip_path = os.path.join(temp_dir, f"{model_name}.zip")
+            
+            # Download the model
+            print(f"Downloading {model_name} model...")
+            try:
+                download_url(url, zip_path)
+            except Exception as e:
+                print(f"Error downloading {model_name}: {e}")
+                print(f"You can manually download it from {url}")
+                continue
+            
+            # Extract the model
+            try:
+                with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+                    zip_ref.extractall(temp_dir)
+                
+                # Move the model file to the destination
+                model_file = f"{model_name}.pt"
+                src_path = os.path.join(temp_dir, model_file)
+                dest_path = os.path.join(destination, model_file)
+                
+                if os.path.exists(src_path):
+                    shutil.move(src_path, dest_path)
+                    print(f"Installed {model_file} to {dest_path}")
+                else:
+                    print(f"Warning: Model file {model_file} not found after extraction")
+            except Exception as e:
+                print(f"Error extracting {model_name}: {e}")
+                print(f"You can manually download and extract it from {url}")
+    
+    print("\nDownload complete!")
+    print(f"Models are installed in: {destination}")
+
+
+def check_models_exist():
+    """Check if models exist in the package directory."""
+    model_dir = get_checkpoint_dir()
+    
+    if not os.path.exists(model_dir):
+        return False
+    
+    # Check if all model files exist
+    all_exist = True
+    for model_name in MODEL_URLS.keys():
+        model_path = os.path.join(model_dir, f"{model_name}.pt")
+        if not os.path.exists(model_path):
+            all_exist = False
+            break
+    
+    return all_exist
+
+
+def get_model_path(model_name):
+    """Get the path to a specific model file."""
+    model_dir = get_checkpoint_dir()
+    return os.path.join(model_dir, f"{model_name}.pt")
+
+
+def print_checkpoint_paths():
+    """Print the checkpoint directory path and available models."""
+    checkpoint_dir = get_checkpoint_dir()
+    print(f"Checkpoint directory: {checkpoint_dir}")
+    
+    if os.path.exists(checkpoint_dir):
+        print("\nAvailable models:")
+        for model_name in MODEL_URLS.keys():
+            model_path = os.path.join(checkpoint_dir, f"{model_name}.pt")
+            if os.path.exists(model_path):
+                print(f"  [+] {model_name}.pt")
+            else:
+                print(f"  [-] {model_name}.pt (not downloaded)")
+    else:
+        print("\nCheckpoint directory does not exist. No models are downloaded.")
+        print("Run 'msfiddle-download-models' to download the models.")
+
+
+# Command-line interface for the download_models function
+def download_models_cli():
+    parser = argparse.ArgumentParser(description="Download pre-trained models for msfiddle")
+    parser.add_argument("-d", "--destination", help="Path to save the models (optional, defaults to package directory)")
+    parser.add_argument("-m", "--models", nargs="+", help="Specific model names to download (optional, defaults to all models)")
+    args = parser.parse_args()
+    
+    # If destination is not provided, it will be None and download_models will use get_checkpoint_dir()
+    download_models(destination=args.destination, models=args.models)
+
+
+# Command-line interface for the print_checkpoint_paths function
+def print_checkpoint_paths_cli():
+    parser = argparse.ArgumentParser(description="Print checkpoint paths and available models")
+    parser.parse_args()  # This just consumes any arguments passed
+    print_checkpoint_paths()
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python download.py [command]")
+        print("Commands:")
+        print("  print_checkpoint_paths - Print the checkpoint directory path and available models")
+        print("  download_models [destination] - Download pre-trained models to the specified destination")
+        sys.exit(1)
+
+    if sys.argv[1] == "print_checkpoint_paths":
+        print_checkpoint_paths()
+    
+    elif sys.argv[1] == "download_models":
+        # If the first argument is 'download', treat the second as destination
+        if len(sys.argv) > 2:
+            download_models(destination=sys.argv[2])
+        else:
+            download_models()
+    
+    else:
+        raise ValueError(f"Unknown command: {sys.argv[1]}")

msfiddle/main.py

@@ -0,0 +1,356 @@
+import os
+import argparse
+from tqdm import tqdm
+import yaml
+import time
+import sys
+
+import numpy as np
+import pandas as pd
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from .dataset import MGFDataset
+from .model_tcn import MS2FNet_tcn, FDRNet
+from .utils.mol_utils import vector_to_formula, formula_to_vector
+from .utils.msms_utils import mass_calculator
+from .utils.refine_utils import formula_refinement
+from .download import get_checkpoint_dir, check_models_exist, download_models
+
+
+def test_step(model, loader, device): 
+    model.eval()
+    spec_ids = []
+    y_pred = []
+    exp_precursor_mz = []
+    exp_precursor_type = []
+    mass_pred = []
+    atomnum_pred = []
+    hcnum_pred = []
+    with tqdm(total=len(loader)) as bar: 
+        for _, batch in enumerate(loader): 
+            spec_id, exp_pre_type, x, env, neutral_add = batch
+            x = x.to(device, dtype=torch.float32)
+            env = env.to(device, dtype=torch.float32)
+            neutral_add = neutral_add.to(device, dtype=torch.float32)
+            exp_pre_mz = env[:, 0]
+
+            with torch.no_grad(): 
+                _, pred_f, pred_mass, pred_atomnum, pred_hcnum = model(x, env)
+            pred_f = pred_f - neutral_add # add the neutral adduct
+
+            bar.set_description('Eval')
+            bar.update(1)
+
+            spec_ids = spec_ids + list(spec_id)
+            y_pred.append(pred_f.detach().cpu())
+            exp_precursor_mz.append(exp_pre_mz.detach().cpu())
+            exp_precursor_type = exp_precursor_type + list(exp_pre_type)
+            mass_pred.append(pred_mass.detach().cpu())
+            atomnum_pred.append(pred_atomnum.detach().cpu())
+            hcnum_pred.append(pred_hcnum.detach().cpu())
+
+    y_pred = torch.cat(y_pred, dim = 0)
+    exp_precursor_mz = torch.cat(exp_precursor_mz, dim = 0)
+    mass_pred = torch.cat(mass_pred, dim = 0)
+    atomnum_pred = torch.cat(atomnum_pred, dim = 0)
+    hcnum_pred = torch.cat(hcnum_pred, dim = 0)
+    return spec_ids, y_pred, exp_precursor_mz, exp_precursor_type, mass_pred, atomnum_pred, hcnum_pred
+
+
+def rerank_by_fdr(fdr_model, spec, env, refined_results, device, K): 
+    fdr_model.eval()
+
+    refine_f = [f for f in refined_results['formula'] if f != None]
+    refine_m = [m for m in refined_results['mass'] if m != None]
+    if len(refine_f) == 0: 
+        refined_results['fdr'] = [0.] * K
+        return refined_results
+
+    # convert refine_f (formula strings) to f (formula vectors)
+    f = [formula_to_vector(f_str) for f_str in refine_f]
+    f = torch.from_numpy(np.array(f))
+
+    spec = spec.to(device, dtype=torch.float32).repeat(f.size(0), 1)
+    env = env.to(device, dtype=torch.float32).repeat(f.size(0), 1)
+    f = f.to(device, dtype=torch.float32)
+    
+    with torch.no_grad(): 
+        fdr = fdr_model(spec, env, f)
+        fdr = torch.sigmoid(fdr).detach().cpu().numpy()
+    
+    # Create a list of tuples with (fdr_value, refine_f_value, mass_value) pairs
+    combined_list = list(zip(fdr, refine_f, refine_m))
+    # Sort the combined list based on the fdr_value
+    sorted_combined_list = sorted(combined_list, key=lambda x: x[0], reverse=True)
+    # Unpack the sorted lists
+    sorted_fdr, sorted_refine_f, sorted_mass = zip(*sorted_combined_list)
+    sorted_fdr, sorted_refine_f, sorted_mass = list(sorted_fdr), list(sorted_refine_f), list(sorted_mass)
+
+    # Pad sorted_refine_f and sorted_fdr with None if necessary
+    if len(sorted_refine_f) < K: 
+        sorted_refine_f += [None] * (K - len(sorted_refine_f))
+        sorted_fdr += [0.] * (K - len(sorted_fdr))
+        sorted_mass += [None] * (K - len(sorted_mass))
+
+    return {'formula': sorted_refine_f, 'mass': sorted_mass, 'fdr': sorted_fdr}
+
+
+def init_random_seed(seed):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    return
+
+
+def main(): 
+    parser = argparse.ArgumentParser(description='msfiddle')
+    
+    # Define two exclusive argument groups
+    mode_group = parser.add_mutually_exclusive_group(required=True)
+    mode_group.add_argument('--demo', action='store_true',
+                      help='Run prediction on demo data')
+    mode_group.add_argument('--test_data', type=str,
+                      help='Path to data (.mgf)')
+    
+    # Add instrument type option
+    parser.add_argument('--instrument_type', type=str, choices=['qtof', 'orbitrap'], default='orbitrap',
+                      help='Instrument type: qtof or orbitrap (default: orbitrap)')
+    
+    # Add other required arguments
+    parser.add_argument('--result_path', type=str, required=True,
+                      help='Path to save predicted results')
+    
+    # Add optional arguments
+    parser.add_argument('--buddy_path', type=str, default='', 
+                      help='Path to saved BUDDY\'s results')
+    parser.add_argument('--sirius_path', type=str, default='', 
+                      help='Path to saved SIRIUS\'s results')
+    parser.add_argument('--seed', type=int, default=42,
+                      help='Seed for random functions')
+    parser.add_argument('--device', type=int, nargs='+', default=[0], 
+                      help='Which GPUs to use if any (default: [0]). Accepts multiple values separated by space.')
+    parser.add_argument('--no_cuda', action='store_true', 
+                      help='Disables CUDA')
+    
+    # Add advanced arguments for custom configurations
+    advanced_group = parser.add_argument_group('Advanced options')
+    advanced_group.add_argument('--config_path', type=str,
+                      help='Custom path to configuration (.yaml)')
+    advanced_group.add_argument('--resume_path', type=str,
+                      help='Custom path to pretrained TCN model')
+    advanced_group.add_argument('--fdr_resume_path', type=str,
+                      help='Custom path to pretrained FDR model')
+    
+    args = parser.parse_args()
+    
+    # Initialize random seed
+    init_random_seed(args.seed)
+    start_time = time.time()
+    
+    # Get package directory
+    package_dir = os.path.dirname(os.path.abspath(__file__))
+    
+    # Make sure models are downloaded
+    if not check_models_exist():
+        print("Pre-trained models not found. Downloading now...")
+        download_models()
+        print("Models downloaded successfully.")
+    
+    # Set default paths based on mode and instrument type
+    if args.demo:
+        test_data_path = os.path.join(package_dir, 'demo', 'input_msms.mgf')
+        print(f"Using demo data: {test_data_path}")
+    else:
+        test_data_path = args.test_data
+        print(f"Using custom data: {test_data_path}")
+    
+    # Set paths based on instrument type
+    instrument_suffix = args.instrument_type
+    
+    # Set config path
+    if args.config_path:
+        config_path = args.config_path
+    else:
+        config_path = os.path.join(package_dir, 'config', f'fiddle_tcn_{instrument_suffix}.yml')
+    print(f"Using config: {config_path}")
+    
+    # Set model paths
+    checkpoint_dir = get_checkpoint_dir()
+    
+    if args.resume_path:
+        resume_path = args.resume_path
+    else:
+        resume_path = os.path.join(checkpoint_dir, f'fiddle_tcn_{instrument_suffix}.pt')
+    print(f"Using TCN model: {resume_path}")
+    
+    if args.fdr_resume_path:
+        fdr_resume_path = args.fdr_resume_path
+    else:
+        fdr_resume_path = os.path.join(checkpoint_dir, f'fiddle_fdr_{instrument_suffix}.pt')
+    print(f"Using FDR model: {fdr_resume_path}")
+    
+    # Load config
+    with open(config_path, 'r') as f: 
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    print('Loaded model & training configuration from {}'.format(config_path))
+
+    # Set device
+    device_1st = torch.device("cuda:" + str(args.device[0])) if torch.cuda.is_available() and not args.no_cuda else torch.device("cpu")
+    print(f'Device(s): {args.device}')
+
+    # 1. Data
+    valid_set = MGFDataset(test_data_path, config['encoding'])
+    valid_loader = DataLoader(
+                    valid_set,
+                    batch_size=1, 
+                    shuffle=False, 
+                    num_workers=0, 
+                    drop_last=True)
+
+    # 2. Model
+    # 2.1 MS2F
+    model = MS2FNet_tcn(config['model']).to(device_1st)
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f'# MS2FNet_tcn Params: {num_params}')
+    if len(args.device) > 1: # Wrap the model with nn.DataParallel
+        model = nn.DataParallel(model, device_ids=args.device)
+
+    print('Loading the formula prediction model...')
+    state_dict = torch.load(resume_path, map_location=device_1st, weights_only=False)['model_state_dict']
+    is_multi_gpu = any(key.startswith('module.') for key in state_dict.keys())
+    if is_multi_gpu and len(args.device) == 1: # Convert the model to single GPU
+        new_state_dict = OrderedDict()
+        for key, value in state_dict.items():
+            if key.startswith('module.'):
+                new_key = key[7:]  # Remove the 'module.' prefix
+                new_state_dict[new_key] = value
+            else:
+                new_state_dict[key] = value
+        model.load_state_dict(new_state_dict)
+    else:
+        model.load_state_dict(state_dict)
+
+    # 2.2 FDR
+    fdr_model = FDRNet(config['model']).to(device_1st)
+    num_params = sum(p.numel() for p in fdr_model.parameters())
+    print(f'# FDRNet Params: {num_params}')    
+    if len(args.device) > 1: # Wrap the model with nn.DataParallel
+        fdr_model = nn.DataParallel(fdr_model, device_ids=args.device)
+
+    print('Loading the FDR prediction model...')
+    state_dict = torch.load(fdr_resume_path, map_location=device_1st, weights_only=False)['model_state_dict']
+    is_multi_gpu = any(key.startswith('module.') for key in state_dict.keys())
+    if is_multi_gpu and len(args.device) == 1: # Convert the fdr_model to single GPU
+        new_state_dict = OrderedDict()
+        for key, value in state_dict.items(): 
+            if key.startswith('module.'):
+                new_key = key[7:]  # Remove the 'module.' prefix
+                new_state_dict[new_key] = value
+            else:
+                new_state_dict[key] = value
+        fdr_model.load_state_dict(new_state_dict)
+    else:
+        fdr_model.load_state_dict(state_dict)
+    
+    # 3. Formula Prediction
+    spec_ids, y_pred, exp_precursor_mz, exp_precursor_type, mass_pred, atomnum_pred, hcnum_pred = test_step(model, valid_loader, device_1st)
+
+    prediction_time = time.time() - start_time
+    prediction_time /= len(valid_set)
+
+    formula_pred = [vector_to_formula(y) for y in y_pred] # calculate the formula strings
+    y_pred = [';'.join(y) for y in y_pred.numpy().astype('str')]
+
+    spectra = []
+    environments = []
+    for batch in valid_loader:  
+        _, _, spec, env, _ = batch
+        spectra.append(spec)
+        environments.append(env)
+
+    # 4. Post-processing
+    if args.buddy_path != '':
+        buddy_df = pd.read_csv(args.buddy_path)
+    if args.sirius_path != '':
+        sirius_df = pd.read_csv(args.sirius_path)
+
+    formula_redined = {'Refined Formula ({})'.format(str(k)): [] for k in range(config['post_processing']['top_k'])}
+    mass_redined = {'Refined Mass ({})'.format(str(k)): [] for k in range(config['post_processing']['top_k'])}
+    fdr_refined = {'FDR ({})'.format(str(k)): [] for k in range(config['post_processing']['top_k'])}
+    running_time = []
+    exp_mass = []
+    # Please note that here we use the experimental precursor m/z, rather than the theoretic precursor m/z. 
+    for idx, pred_f, exp_pre_mz, exp_pre_type, spec, env in tqdm(zip(spec_ids, 
+                                                             formula_pred, 
+                                                             exp_precursor_mz, 
+                                                             exp_precursor_type, spectra, environments), total=len(exp_precursor_mz), desc='Post'): 
+        m = mass_calculator(exp_pre_type, exp_pre_mz) # Use experimental precursor m/z and precursor type to calculate molmass
+        exp_mass.append(m.item())
+
+        f0_list = [pred_f]
+        if args.buddy_path != '' and len(buddy_df.loc[buddy_df['ID'] == idx]) > 0: 
+            buddy_f = buddy_df.loc[buddy_df['ID'] == idx].iloc[0][['Pred Formula (1)', 
+                                                                        'Pred Formula (2)', 
+                                                                        'Pred Formula (3)',
+                                                                        'Pred Formula (4)', 
+                                                                        'Pred Formula (5)',]].tolist()
+            buddy_fdr = buddy_df.loc[buddy_df['ID'] == idx].iloc[0][['BUDDY Score (1)', 
+                                                                        'BUDDY Score (2)', 
+                                                                        'BUDDY Score (3)',
+                                                                        'BUDDY Score (4)', 
+                                                                        'BUDDY Score (5)',]].tolist()
+            buddy_f = [x for x, fdr in zip(buddy_f, buddy_fdr) if str(x) != 'nan' and fdr < config['post_processing']['buddy_fdr_thr']]
+            f0_list.extend(buddy_f)
+        if args.sirius_path != '': 
+            sirius_f = sirius_df.loc[sirius_df['ID'] == idx].iloc[0][['Pred Formula (1)', 
+                                                                        'Pred Formula (2)', 
+                                                                        'Pred Formula (3)',
+                                                                        'Pred Formula (4)', 
+                                                                        'Pred Formula (5)',]].tolist()
+            sirius_score = sirius_df.loc[sirius_df['ID'] == idx].iloc[0][['SIRIUS Score (1)', 
+                                                                        'SIRIUS Score (2)', 
+                                                                        'SIRIUS Score (3)',
+                                                                        'SIRIUS Score (4)', 
+                                                                        'SIRIUS Score (5)',]].tolist()
+            sirius_f = [x for x, score in zip(sirius_f, sirius_score) if str(x) != 'nan' and score > config['post_processing']['sirius_score_thr']]
+            f0_list.extend(sirius_f)
+
+        f0_list = list(set(f0_list)) # deduplicates
+        start_time = time.time()
+        refined_results = formula_refinement(f0_list, m.item(), 
+                                            config['post_processing']['mass_tolerance'], 
+                                            config['post_processing']['ppm_mode'], 
+                                            config['post_processing']['top_k'], 
+                                            config['post_processing']['maxium_miss_atom_num'], 
+                                            config['post_processing']['time_out'], 
+                                            config['post_processing']['refine_atom_type'],
+                                            config['post_processing']['refine_atom_num'],
+                                            )
+
+        # Rerank the results by predicted FDR
+        refined_results = rerank_by_fdr(fdr_model, spec, env, refined_results, device_1st, config['post_processing']['top_k'])
+
+        for i, (refined_f, refined_m, refined_fdr) in enumerate(zip(refined_results['formula'], refined_results['mass'], refined_results['fdr'])): 
+            formula_redined[f'Refined Formula ({i})'].append(refined_f)
+            mass_redined[f'Refined Mass ({i})'].append(refined_m)
+            fdr_refined[f'FDR ({i})'].append(refined_fdr)
+        refinement_time = time.time() - start_time
+        running_time.append(prediction_time + refinement_time)
+
+    # 5. Save the final results
+    print('\nSaving predicted results...')
+    out_dict = {'ID': spec_ids, 'Y Pred': y_pred, 'Mass': exp_mass, 
+                'Pred Formula': formula_pred, 'Pred Mass': mass_pred.tolist(), 
+                'Pred Atom Num': atomnum_pred.tolist(), 'Pred H/C Num': hcnum_pred.tolist(), 
+                'Running Time': running_time}
+    res_df = pd.DataFrame({**out_dict, **formula_redined, **mass_redined, **fdr_refined})
+    res_df.to_csv(args.result_path, index=False)
+    print(f'Done! Results saved to {args.result_path}')
+
+
+if __name__ == "__main__":
+    main()