macOS / Apple MPS support, dependency fixes, and faster inference

.gitignore

@@ -0,0 +1,10 @@
+__pycache__/
+*.pyc
+.venv/
+/data/
+cascadia_*/
+checkpoint_*/
+logs/
+*.egg-info/
+build/
+dist/

cascadia/augment.py

@@ -133,27 +133,32 @@ def write_asf(outfile, prec_to_spec, scan_width, max_pep_length, max_charge):
         ms1_rts = ms1_rts[sorted_ms1_rt_idxs]
         ms1_scans = ms1_scans[sorted_ms1_rt_idxs]
 
-        for charge in range(1,max_charge+1):
-            count += 1
+        # The peak list is identical for every candidate charge state, so
+        # format it once and reuse it.
+        peak_lines = []
+        for scan, cur_rt in zip(scans,rts):
+            for mz, intensity in scan:
+                peak_lines.append(f"{mz}\t{intensity}\t{cur_rt}\t2\n")
 
-            out.write("BEGIN IONS\n")
-            out.write(f"TITLE={str(count)}\n")
-            out.write(f"PEPMASS={prec}\n")
-            out.write(f"CHARGE={charge}\n")
-            out.write(f"SCAN={count}\n")
-            out.write(f"RT={rt}\n")
-            out.write(f"SEQ={'K'*max_pep_length}\n")
-
-            for scan, cur_rt in zip(scans,rts):
-                for mz, intensity in scan:
-                    out.write(f"{mz}\t{intensity}\t{cur_rt}\t2\n")
-
-            for scan, cur_rt in zip(ms1_scans, ms1_rts):
-                for mz, intensity in scan:
-                    if np.abs(mz - prec) < window_width + 1:
-                        out.write(f"{mz}\t{intensity}\t{cur_rt}\t1\n")
+        for scan, cur_rt in zip(ms1_scans, ms1_rts):
+            for mz, intensity in scan:
+                if np.abs(mz - prec) < window_width + 1:
+                    peak_lines.append(f"{mz}\t{intensity}\t{cur_rt}\t1\n")
+        peak_block = ''.join(peak_lines)
 
-            out.write("END IONS\n")   
+        for charge in range(1,max_charge+1):
+            count += 1
+            out.write(
+                "BEGIN IONS\n"
+                f"TITLE={str(count)}\n"
+                f"PEPMASS={prec}\n"
+                f"CHARGE={charge}\n"
+                f"SCAN={count}\n"
+                f"RT={rt}\n"
+                f"SEQ={'K'*max_pep_length}\n"
+                + peak_block +
+                "END IONS\n"
+            )
     out.close()
 
 def augment_spectra(mzml_file, temp_path, top_n = 150, scan_width = 1, max_pep_length = 30, max_charge = 3):

cascadia/cascadia.py

@@ -3,30 +3,98 @@
 from .depthcharge.tokenizers import PeptideTokenizer
 import torch
 import numpy as np
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 import os
 import sys
 import argparse
 from lightning.pytorch import loggers as pl_loggers
 from .depthcharge.utils import *
 from .model import AugmentedSpec2Pep
 from .augment import *
+from .utils import write_results
 from datetime import datetime
+import time
 import warnings
 import json
 warnings.filterwarnings('ignore')
 
+MODEL_HPARAMS = dict(
+    d_model=512,
+    n_layers=9,
+    n_head=8,
+    dim_feedforward=1024,
+    dropout=0,
+)
+
+
+def select_accelerator(requested='auto'):
+  """Pick the Lightning accelerator: CUDA, then Apple MPS, then CPU.
+
+  ``requested`` may be 'auto', 'gpu' (CUDA), 'mps', or 'cpu'.
+  """
+  if requested != 'auto':
+    print(f'Using requested device: {requested}')
+    return requested
+  if torch.cuda.is_available():
+    print('GPU found')
+    return 'gpu'
+  if torch.backends.mps.is_available():
+    print('Apple MPS GPU found')
+    return 'mps'
+  print('No GPU found - running on cpu')
+  return 'cpu'
+
+
+def dataloader_workers(n_workers):
+  """Number of DataLoader worker processes to use.
+
+  macOS and Windows start worker processes with ``spawn``, which requires
+  pickling the dataset; the tokenizer's numba typed dictionaries are not
+  picklable, so spectra are loaded in the main process there. Linux forks
+  the workers instead. Reading batches from the HDF5 index is cheap compared
+  to inference, so this costs little.
+  """
+  return n_workers if sys.platform.startswith('linux') else 0
+
+
+def load_tokenizer(mods):
+  if mods == 'mskb':
+    return PeptideTokenizer.from_massivekb(reverse=False, replace_isoleucine_with_leucine=True)
+  with open(mods, 'r') as f:
+    proforma = json.load(f)
+  return PeptideTokenizer.from_proforma(proforma, reverse=False, replace_isoleucine_with_leucine=True)
+
+
+def load_model(model_ckpt_path, tokenizer, rt_width, max_charge, lr=1e-4):
+  """Build an AugmentedSpec2Pep and load weights from a Lightning checkpoint.
+
+  Cascadia checkpoints contain numpy scalars in their optimizer/scheduler
+  state, which the ``weights_only=True`` default of ``torch.load`` (PyTorch
+  >= 2.6) refuses to unpickle, so the weights are loaded explicitly here.
+  """
+  checkpoint = torch.load(model_ckpt_path, map_location='cpu', weights_only=False)
+  model = AugmentedSpec2Pep(
+      rt_width=rt_width,
+      tokenizer=tokenizer,
+      max_charge=max_charge,
+      lr=lr,
+      **MODEL_HPARAMS,
+  )
+  model.load_state_dict(checkpoint['state_dict'])
+  return model
+
 
 def sequence():
   parser=argparse.ArgumentParser()
   parser.add_argument("spectrum_file")
   parser.add_argument("model", type= str, help="A path to a trained Cascadia model checkpoint.")
   parser.add_argument("-o", "--outfile", type= str, default= 'cascadia_results', help="Output file for inference.")
   parser.add_argument("-t", "--score_threshold", type= float, default= 0.8, help="Score threshold for Cascadia predictions.")
-  parser.add_argument("-b", "--batch_size", type= int, default= 32, help="Number of spectra to include in a batch.")
+  parser.add_argument("-b", "--batch_size", type= int, default= 128, help="Number of spectra to include in a batch. Larger batches are faster on a GPU as long as they fit in memory.")
   parser.add_argument("-w", "--width", type= int, default= 2, help="Number of adjacent scans to use when constructing each augmented spectrum.")
   parser.add_argument("-c", "--max_charge", type= int, default= 4, help="Maximum precursor charge state to consider")
   parser.add_argument("-p", "--modifications", type= str, default= 'mskb', help="A path to a json file containing a list of the PTMs to consider.")
+  parser.add_argument("-d", "--device", type= str, default= 'auto', choices=['auto', 'gpu', 'mps', 'cpu'], help="Device to run inference on. 'auto' picks CUDA, then Apple MPS, then CPU.")
 
   args = parser.parse_args(args=sys.argv[2:])
 
@@ -43,51 +111,42 @@ def sequence():
   os.mkdir(temp_path)
   train_index_filename = temp_path + "/index.hdf5"
 
+  t_start = time.perf_counter()
   print("Augmenting spectra from:", spectrum_file)
   asf_file, isolation_window_size, cycle_time = augment_spectra(spectrum_file, temp_path, max_charge=max_charge)
-
-  if mods == 'mskb':
-    tokenizer = PeptideTokenizer.from_massivekb(reverse=False, replace_isoleucine_with_leucine=True)
-  else:
-    with open('ptms.json', 'r') as f:
-      proforma = json.load(f)
-    tokenizer = PeptideTokenizer.from_proforma(proforma, reverse=False, replace_isoleucine_with_leucine=True)
-
+  t_augment = time.perf_counter()
+
+  tokenizer = load_tokenizer(mods)
+
+  device = select_accelerator(args.device)
+
   train_dataset = AnnotatedSpectrumDataset(tokenizer, asf_file, index_path=train_index_filename, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm])
-  train_loader = train_dataset.loader(batch_size=batch_size, num_workers=4, pin_memory=True)
+  train_loader = train_dataset.loader(batch_size=batch_size, num_workers=dataloader_workers(4), pin_memory=(device == 'gpu'))
 
   if os.path.exists(asf_file):
       os.remove(asf_file)
+  t_index = time.perf_counter()
 
-  model = AugmentedSpec2Pep.load_from_checkpoint(
-      model_ckpt_path,
-      d_model = 512,
-      n_layers = 9,
-      n_head = 8,
-      dim_feedforward = 1024,
-      dropout = 0,
-      rt_width = 2,
-      tokenizer=tokenizer,
-      max_charge=10,
-  )
+  model = load_model(model_ckpt_path, tokenizer, rt_width=2, max_charge=10)
 
   print("Running inference on augmented spectra from:", spectrum_file)
-  if torch.cuda.is_available():
-    device = 'gpu'
-    print('GPU found')
-  else:
-    device = 'cpu'
-    print(f'No GPU found - running inference on cpu')
 
-  trainer = pl.Trainer(max_epochs=50, log_every_n_steps=1, accelerator=device, devices = 1)
+  trainer = pl.Trainer(max_epochs=50, log_every_n_steps=1, accelerator=device, devices = 1, logger=False)
   preds = trainer.predict(model, dataloaders=train_loader)
+  t_predict = time.perf_counter()
 
   print("Writing results to:", results_file + '.ssl')
 
   write_results(preds, results_file, spectrum_file, isolation_window_size, score_threshold, augmentation_width*cycle_time)
 
   os.remove(train_index_filename)
   os.rmdir(temp_path)
+  t_end = time.perf_counter()
+  print(
+      f"Done in {t_end - t_start:.1f}s "
+      f"(augmentation {t_augment - t_start:.1f}s, indexing {t_index - t_augment:.1f}s, "
+      f"model loading + inference {t_predict - t_index:.1f}s, writing results {t_end - t_predict:.1f}s)"
+  )
 
   return parser
 
@@ -116,11 +175,8 @@ def train():
   lr = args.learning_rate
   mods = args.modifications
 
-  if torch.cuda.is_available():
-    device = 'gpu'
-    print('GPU found')
-  else:
-    device = 'cpu'
+  device = select_accelerator()
+  if device == 'cpu':
     print("No GPU Available - training on CPU will be extremely slow!")
 
   print("Training on spectra from:", train_spectrum_file)
@@ -136,51 +192,31 @@ def train():
   if os.path.exists(val_index_filename):
     os.remove(val_index_filename)
 
-  if mods == 'mskb':
-    tokenizer = PeptideTokenizer.from_massivekb(reverse=False, replace_isoleucine_with_leucine=True)
-  else:
-    with open(mods, 'r') as f:
-      proforma = json.load(f)
-    tokenizer = PeptideTokenizer.from_proforma(proforma, reverse=False, replace_isoleucine_with_leucine=True)
-
+  tokenizer = load_tokenizer(mods)
+
   if '.hdf5' in train_spectrum_file:
     train_dataset = AnnotatedSpectrumDataset(tokenizer, index_path=train_spectrum_file, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm])
     val_dataset = AnnotatedSpectrumDataset(tokenizer, index_path=val_spectrum_file, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm])
   else:
     train_dataset = AnnotatedSpectrumDataset(tokenizer, train_spectrum_file, index_path=train_index_filename, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm]) 
     val_dataset = AnnotatedSpectrumDataset(tokenizer, val_spectrum_file, index_path=val_index_filename, preprocessing_fn=[scale_intensity(scaling="root"), scale_to_unit_norm])
 
-  train_loader = train_dataset.loader(batch_size=batch_size, num_workers=10, pin_memory=True, shuffle=True)
-  val_loader = val_dataset.loader(batch_size=batch_size, num_workers=10, pin_memory=True)
+  train_loader = train_dataset.loader(batch_size=batch_size, num_workers=dataloader_workers(10), pin_memory=(device == 'gpu'), shuffle=True)
+  val_loader = val_dataset.loader(batch_size=batch_size, num_workers=dataloader_workers(10), pin_memory=(device == 'gpu'))
 
   if model_ckpt_path is None:
     print("Training model from scratch") 
     model = AugmentedSpec2Pep(
-        d_model = 512,
-        n_layers = 9,
-        n_head = 8,
-        dim_feedforward = 1024,
-        dropout = 0,
         rt_width = augmentation_width,
         tokenizer=tokenizer,
         max_charge=max_charge,
-        lr=lr
+        lr=lr,
+        **MODEL_HPARAMS,
     )
 
   else:
     print("Loading model from checkpoint:", model_ckpt_path) 
-    model = AugmentedSpec2Pep.load_from_checkpoint(
-      model_ckpt_path, 
-      d_model = 512,
-      n_layers = 9,
-      n_head = 8,
-      dim_feedforward = 1024,
-      dropout = 0,
-      rt_width = 2,
-      tokenizer=tokenizer,
-      max_charge=max_charge,
-      lr=lr
-    )
+    model = load_model(model_ckpt_path, tokenizer, rt_width=2, max_charge=max_charge, lr=lr)
 
   tb_logger = pl_loggers.TensorBoardLogger(save_dir="logs/")
   ckpt_callback = pl.callbacks.ModelCheckpoint(