support the selection of a GPU to be used.

README.md

@@ -56,6 +56,21 @@ It may also be compatible with other formats of spectral libraries based on requ
 - `-ws`<br>
 This parameter is used to specify the .d folder or the folder containing multiple .d folders that need to be analyzed.
 
+Other optional params are list below by entering `beta_dia -h`:
+```bash
+       ******************
+       * Beta-DIA x.y.z *
+       ******************
+Usage: beta_dia [-h] -ws WS -lib LIB [-out_name OUT_NAME] [-gpu_id GPU_ID]
+
+optional arguments:
+  -h, --help          show this help message and exit
+  -ws WS              specify the folder that is .d or contains .d files.
+  -lib LIB            specify the absolute path of a .speclib spectra library.
+  -out_name OUT_NAME  specify the folder name of outputs. Default: beta_dia.
+  -gpu_id GPU_ID      specify the GPU-ID (e.g. 0, 1, 2) which will be used. Default: 0.
+```
+
 ### Output
 Beta-DIA will generate **`beta_dia/report_beta.log.txt`** and **`beta_dia/report_beta.tsv`** in each .d folder. 
 The report_beta.tsv contains precursor and protein IDs, as well as plenty of associated information. 
@@ -92,5 +107,8 @@ optimize the proteome profiling of diaPASEF mass spectrometry data**
 ## Changelog
 
 ### 0.1.0
-
   * FEAT: first commit on GitHub.
+### 0.1.1
+  * FIX: loading .d even with missing or occupied frame.
+### 0.1.2
+  * FEAT: support the selection of a GPU to be used.

beta_dia/__init__.py

@@ -1 +1 @@
-__version__ = '0.1.1'
+__version__ = '0.1.2'

beta_dia/deepmall.py

@@ -57,7 +57,7 @@ def extract_mall(
     ims = ims.permute((0, 2, 1))
     ims = ims[:, (center_idx - 1) : (center_idx + 2), :]
     pred_ims = df_batch['pred_im'].values
-    pred_ims = torch.from_numpy(pred_ims).cuda()
+    pred_ims = torch.from_numpy(pred_ims).to(param_g.gpu_id)
     pred_ims = pred_ims.unsqueeze(-1).unsqueeze(-1).expand(ims.shape)
     bias_ims = pred_ims - ims
     bias_ims[ims < 0] = param_g.tol_im_xic
@@ -67,7 +67,7 @@ def extract_mall(
     mzs = mzs.permute((0, 2, 1))
     mzs = mzs[:, (center_idx - 1) : (center_idx + 2), :]
     pred_mzs = np.stack(df_batch['fg_mz'])
-    pred_mzs = torch.from_numpy(pred_mzs).cuda()
+    pred_mzs = torch.from_numpy(pred_mzs).to(param_g.gpu_id)
     pred_mzs = pred_mzs.unsqueeze(1).expand(mzs.shape)
     ppms = 1e6 * (pred_mzs - mzs) / (pred_mzs + 1e-7)
     ppms[mzs < 1] = param_g.tol_ppm
@@ -76,7 +76,7 @@ def extract_mall(
     # sa
     cols = ['score_center_elution_' + str(i) for i in range(14)]
     elutions = df_batch[cols].values[:, 2:]
-    elutions = torch.from_numpy(elutions).cuda()
+    elutions = torch.from_numpy(elutions).to(param_g.gpu_id)
     elutions = elutions.unsqueeze(1)
 
     # area
@@ -89,23 +89,23 @@ def extract_mall(
     rts = np.repeat(rts[:, np.newaxis, :], xics.shape[1], axis=1)
     areas = np.trapz(xics, x=rts, axis=2)
     areas = areas / (areas.max(axis=1, keepdims=True) + 1e-7)
-    areas = torch.from_numpy(areas).cuda()
+    areas = torch.from_numpy(areas).to(param_g.gpu_id)
     areas = areas.unsqueeze(1)
 
     # pred intensities
     pred_heights = np.stack(df_batch['fg_height'])
-    pred_heights = torch.from_numpy(pred_heights).cuda()
+    pred_heights = torch.from_numpy(pred_heights).to(param_g.gpu_id)
     pred_heights = pred_heights.unsqueeze(1)
 
     # ion type
     fg_type = np.stack(df_batch['fg_anno']) // 1000
-    fg_type = torch.from_numpy(fg_type.astype(np.float32)).cuda()
+    fg_type = torch.from_numpy(fg_type.astype(np.float32)).to(param_g.gpu_id)
     fg_type = fg_type.unsqueeze(1)
 
     # snr
     cols = ['score_center_snr_' + str(i) for i in range(14)]
     snr = df_batch[cols].values[:, 2:]
-    snr = torch.from_numpy(snr).cuda()
+    snr = torch.from_numpy(snr).to(param_g.gpu_id)
     snr = snr.unsqueeze(1)
 
     mall = torch.cat([pred_heights,
@@ -146,7 +146,7 @@ def scoring_mall(
                         tol_im,
                         tol_ppm)
     valid_ion_nums = df_input['fg_num'].values
-    valid_ion_nums = torch.tensor(valid_ion_nums).long().cuda()
+    valid_ion_nums = torch.tensor(valid_ion_nums).long().to(param_g.gpu_id)
     with torch.no_grad():
         feature, pred = model_mall(mall, valid_ion_nums)
 

beta_dia/deepmap.py

@@ -32,7 +32,7 @@ def load_models(dir_center=None, dir_big=None):
 
 def load_model_big(dir_model, channels):
     model = models.DeepMap(channels)
-    device = torch.device('cuda')
+    device = param_g.gpu_id
     if dir_model is None:
         pt_path = Path(__file__).resolve().parent/'pretrained'/'deepbig.pt'
         model.load_state_dict(torch.load(pt_path, map_location=device))
@@ -45,7 +45,7 @@ def load_model_big(dir_model, channels):
 
 def load_model_center(dir_model, channels):
     model = models.DeepMap(channels)
-    device = torch.device('cuda')
+    device = param_g.gpu_id
     if dir_model is None:
         pt_path = Path(__file__).resolve().parent/'pretrained'/'deepcenter.pt'
         model.load_state_dict(torch.load(pt_path, map_location=device))
@@ -442,7 +442,7 @@ def scoring_maps(
         non_fg_num = maps.shape[1] - param_g.fg_num
         valid_ion_nums = non_fg_num + df_batch['fg_num'].values
         valid_ion_nums = torch.tensor(
-            np.repeat(valid_ion_nums, locus_num)).long().cuda()
+            np.repeat(valid_ion_nums, locus_num)).long().to(param_g.gpu_id)
         with torch.no_grad():
             feature, pred = model(maps, valid_ion_nums)
         torch.cuda.synchronize()  # for profile
@@ -565,7 +565,7 @@ def extract_scoring_big(
             valid_ion_nums = 4 * (2 + df_input['fg_num'].values)
             model = model_big
         maps_sub = maps[:, idx]
-        valid_ion_nums = torch.tensor(valid_ion_nums).long().cuda()
+        valid_ion_nums = torch.tensor(valid_ion_nums).long().to(param_g.gpu_id)
         with torch.no_grad():
             feature, pred = model(maps_sub, valid_ion_nums)
         pred = torch.softmax(pred, 1)

beta_dia/dist/pyarmor_runtime_000000/__init__.py

@@ -1,4 +1,4 @@
-# Pyarmor 8.5.11 (trial), 000000, 2024-09-04T18:13:31.585123
+# Pyarmor 8.5.11 (trial), 000000, 2024-09-06T18:03:46.086765
 def __pyarmor__():
     import platform
     import sys

beta_dia/fxic.py

@@ -122,7 +122,7 @@ def cal_coelution_by_gaussion(xics, window_points, valids_num):
     '''
     Coelution scores by sliding windows methods
     '''
-    valids_num = torch.tensor(valids_num, device=param_g.device)
+    valids_num = torch.tensor(valids_num, device=param_g.gpu_id)
 
     # block -- profile
     block_num = xics.shape[0] * xics.shape[1]

beta_dia/param_g.py

@@ -6,10 +6,6 @@
 
 is_save_final = True
 
-device = torch.device('cuda:0')
-device_name = torch.cuda.get_device_name(device)
-torch.backends.cudnn.benchmark = True
-
 # placeholder
 ws = None
 dir_out = None
@@ -30,17 +26,6 @@
 rubbish_q_cut = 0.5
 # protein inference q cut
 inference_q_cut = 0.05
-# xic extraction occupied the GPU memory ratio
-if '4090' in device_name:
-    batch_xic_seed = 5000
-    batch_xic_locus = batch_xic_seed * 5
-    batch_deep_center = 10000
-    batch_deep_big = 5000
-else:
-    batch_xic_seed = 4000
-    batch_xic_locus = batch_xic_seed * 5
-    batch_deep_center = 10000
-    batch_deep_big = 2000
 
 # widely used
 fg_num = 12

beta_dia/refine.py

@@ -215,7 +215,7 @@ def my_collate(items):
 
 
 def eval_one_epoch(trainloader, model):
-    device = param_g.device
+    device = param_g.gpu_id
     model.eval()
     prob_v, label_v = [], []
 
@@ -247,7 +247,7 @@ def eval_one_epoch(trainloader, model):
 
 
 def train_one_epoch(trainloader, model, optimizer, loss_fn):
-    device = param_g.device
+    device = param_g.gpu_id
     model.train()
     epoch_loss = 0.
     for batch_idx, (batch_map, batch_map_len, batch_y) in enumerate(
@@ -360,7 +360,7 @@ def train_model_mall(malls, valid_num, labels, epochs):
 
     # model
     model = models.DeepMall(input_dim=mall_dim,
-                            feature_dim=32).to(param_g.device)
+                            feature_dim=32).to(param_g.gpu_id)
 
     # optimizer
     optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)

beta_dia/scoring.py

@@ -265,7 +265,7 @@ def scoring_xic_intensity(df_batch, xics, rts):
     elutions = df_batch[cols].values + 1e-7
 
     # boundary
-    sa_m = torch.from_numpy(elutions).cuda()
+    sa_m = torch.from_numpy(elutions).to(param_g.gpu_id)
     locus_start_v, locus_end_v = fxic.estimate_xic_boundary(xics, sa_m)
     locus_start_v = locus_start_v.astype(np.int8)
     locus_end_v = locus_end_v.astype(np.int8)

beta_dia/utils.py

@@ -48,7 +48,7 @@ def convert_numba_to_tensor(x):
 
 
 def create_cuda_zeros(shape, dtype=torch.float32):
-    x = torch.zeros(shape, dtype=dtype, device=param_g.device)
+    x = torch.zeros(shape, dtype=dtype, device=param_g.gpu_id)
     x = cuda.as_cuda_array(x)
     return x
 
@@ -276,23 +276,57 @@ def get_args():
     print(' ' * 9, f"* {name} *")
     print(' ' * 9, "*" * (len(name) + 4))
 
-    parser = argparse.ArgumentParser('Beta-DIA for diaPASEF analysis')
+    parser = argparse.ArgumentParser('beta_dia')
+
+    # required=True
     parser.add_argument(
-        '-ws', '--ws', required=True,
+        '-ws', required=True,
         help='specify the folder that is .d or contains .d files.'
     )
     parser.add_argument(
-        '-lib', '--lib', required=True,
+        '-lib', required=True,
         help='specify the absolute path of a .speclib spectra library.'
     )
+
+    # optional
+    parser.add_argument(
+        '-out_name', type=str, default='beta_dia',
+        help='specify the folder name of outputs. Default: beta_dia.'
+    )
     parser.add_argument(
-        '-out_name', '--out_name', type=str, default='beta_dia',
-        help='specify the folder name of outputs.'
+        '-gpu_id', type=int, default=0,
+        help='specify the GPU-ID (e.g. 0, 1, 2) which will be used. Default: 0'
     )
+
     args = parser.parse_args()
+
+    # process params
+    init_gpu_params(args.gpu_id)
+
     return Path(args.ws), Path(args.lib), args.out_name
 
 
+def init_gpu_params(gpu_id):
+    param_g.gpu_id = torch.device('cuda:' + str(gpu_id))
+    param_g.device_name = torch.cuda.get_device_name(gpu_id)
+    torch.backends.cudnn.benchmark = True
+
+    from numba import cuda
+    cuda.select_device(gpu_id)
+
+    # xic extraction occupied the GPU memory ratio
+    if '4090' in param_g.device_name:
+        param_g.batch_xic_seed = 5000
+        param_g.batch_xic_locus = param_g.batch_xic_seed * 5
+        param_g.batch_deep_center = 10000
+        param_g.batch_deep_big = 5000
+    else:
+        param_g.batch_xic_seed = 4000
+        param_g.batch_xic_locus = param_g.batch_xic_seed * 5
+        param_g.batch_deep_center = 10000
+        param_g.batch_deep_big = 2000
+
+
 def init_multi_ws(ws):
     # GPU memory has to be larger than 10G!
     total_memory = torch.cuda.get_device_properties(0).total_memory