geney 1.3.78__py2.py3-none-any.whl → 1.4.0__py2.py3-none-any.whl

geney/utils/pangolin_utils.py

@@ -0,0 +1,173 @@
+# Load models
+#
+__all__ = ['pangolin_predict_probs']
+ # Load models
+import torch
+from pkg_resources import resource_filename
+from pangolin.model import *
+import numpy as np
+import sys
+
+pang_model_nums = [0, 1, 2, 3, 4, 5, 6, 7]
+pang_models = []
+
+device = torch.device('cpu')
+if sys.platform == 'darwin':
+    device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
+
+if sys.platform == 'linux':
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+print(f"Pangolin loaded to {device}.")
+
+for i in pang_model_nums:
+    for j in range(1, 6):
+        model = Pangolin(L, W, AR).to(device)
+        if torch.cuda.is_available():
+            model.cuda()
+            # weights = torch.load(resource_filename("pangolin","models/final.%s.%s.3" % (j, i)))
+            weights = torch.load(resource_filename("pangolin", "models/final.%s.%s.3" % (j, i)), weights_only=True)
+
+        else:
+            weights = torch.load(resource_filename("pangolin","models/final.%s.%s.3" % (j, i)), weights_only=True,
+                                 map_location=device)
+
+        model.load_state_dict(weights)
+        model.eval()
+        pang_models.append(model)
+
+
+def pang_one_hot_encode(seq):
+    IN_MAP = np.asarray([[0, 0, 0, 0],
+                         [1, 0, 0, 0],
+                         [0, 1, 0, 0],
+                         [0, 0, 1, 0],
+                         [0, 0, 0, 1]])
+    seq = seq.upper().replace('A', '1').replace('C', '2')
+    seq = seq.replace('G', '3').replace('T', '4').replace('N', '0')
+    seq = np.asarray(list(map(int, list(seq))))
+    return IN_MAP[seq.astype('int8')]
+
+
+
+def pangolin_predict_probs(true_seq, models, just_ss=False):
+    # print(f"Running pangolin on: {true_seq}")
+    if just_ss:
+        model_nums = [0, 2, 4, 6]
+    else:
+        model_nums = [0, 1, 2, 3, 4, 5, 6, 7]
+
+    INDEX_MAP = {0: 1, 1: 2, 2: 4, 3: 5, 4: 7, 5: 8, 6: 10, 7: 11}
+
+    seq = true_seq
+    true_seq = true_seq[5000:-5000]
+    acceptor_dinucleotide = np.array([true_seq[i - 2:i] == 'AG' for i in range(len(true_seq))]) # np.ones(len(true_seq)) #
+    donor_dinucleotide = np.array([true_seq[i+1:i+3] == 'GT' for i in range(len(true_seq))]) #np.ones(len(true_seq)) #
+
+    seq = pang_one_hot_encode(seq).T
+    seq = torch.from_numpy(np.expand_dims(seq, axis=0)).float()
+
+    # if torch.cuda.is_available():
+    seq = seq.to(torch.device(device))
+
+    scores = []
+    for j, model_num in enumerate(model_nums):
+        score = []
+        # Average across 5 models
+        for model in models[5 * j:5 * j + 5]:
+            with torch.no_grad():
+                score.append(model(seq.to(device))[0][INDEX_MAP[model_num], :].cpu().numpy())
+
+        scores.append(np.mean(score, axis=0))
+
+    splicing_pred = np.array(scores).max(axis=0)
+    donor_probs = [splicing_pred[i] * donor_dinucleotide[i] for i in range(len(true_seq))]
+    acceptor_probs = [splicing_pred[i] * acceptor_dinucleotide[i] for i in range(len(true_seq))]
+    return donor_probs, acceptor_probs
+
+
+#
+# import torch
+# from pkg_resources import resource_filename
+# from pangolin.model import *
+# import numpy as np
+# import sys
+#
+# pang_model_nums = [0, 1, 2, 3, 4, 5, 6, 7]
+# pang_models = []
+#
+# device = torch.device('cpu')
+# if sys.platform == 'darwin':
+#     device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
+#
+# if sys.platform == 'linux':
+#     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+#
+# device = 'cpu'
+# print(f"Pangolin loaded to {device}.")
+#
+# for i in pang_model_nums:
+#     for j in range(1, 6):
+#         model = Pangolin(L, W, AR).to(device)
+#         if torch.cuda.is_available():
+#             model.cuda()
+#             # weights = torch.load(resource_filename("pangolin","models/final.%s.%s.3" % (j, i)))
+#             weights = torch.load(resource_filename("pangolin", "models/final.%s.%s.3" % (j, i)), weights_only=True)
+#
+#         else:
+#             weights = torch.load(resource_filename("pangolin","models/final.%s.%s.3" % (j, i)), weights_only=True,
+#                                  map_location=device)
+#
+#         model.load_state_dict(weights)
+#         model.eval()
+#         pang_models.append(model)
+#
+#
+# def pang_one_hot_encode(seq):
+#     IN_MAP = np.asarray([[0, 0, 0, 0],
+#                          [1, 0, 0, 0],
+#                          [0, 1, 0, 0],
+#                          [0, 0, 1, 0],
+#                          [0, 0, 0, 1]])
+#     seq = seq.upper().replace('A', '1').replace('C', '2')
+#     seq = seq.replace('G', '3').replace('T', '4').replace('N', '0')
+#     seq = np.asarray(list(map(int, list(seq))))
+#     return IN_MAP[seq.astype('int8')]
+#
+#
+#
+# def pangolin_predict_probs(true_seq, models, just_ss=False):
+#     # print(f"Running pangolin on: {true_seq}")
+#     if just_ss:
+#         model_nums = [0, 2, 4, 6]
+#     else:
+#         model_nums = [0, 1, 2, 3, 4, 5, 6, 7]
+#
+#     INDEX_MAP = {0: 1, 1: 2, 2: 4, 3: 5, 4: 7, 5: 8, 6: 10, 7: 11}
+#
+#     seq = true_seq
+#     true_seq = true_seq[5000:-5000]
+#     acceptor_dinucleotide = np.array([true_seq[i - 2:i] == 'AG' for i in range(len(true_seq))]) # np.ones(len(true_seq)) #
+#     donor_dinucleotide = np.array([true_seq[i+1:i+3] == 'GT' for i in range(len(true_seq))]) #np.ones(len(true_seq)) #
+#
+#     seq = pang_one_hot_encode(seq).T
+#     seq = torch.from_numpy(np.expand_dims(seq, axis=0)).float()
+#
+#     # if torch.cuda.is_available():
+#     seq = seq.to(torch.device(device))
+#     print(seq)
+#     scores = []
+#     for j, model_num in enumerate(model_nums):
+#         score = []
+#         # Average across 5 models
+#         for model in models[5 * j:5 * j + 5]:
+#             with torch.no_grad():
+#                 score.append(model(seq)[0][INDEX_MAP[model_num], :].cpu().numpy())
+#
+#         scores.append(np.mean(score, axis=0))
+#
+#     splicing_pred = np.array(scores).max(axis=0)
+#     donor_probs = [splicing_pred[i] * donor_dinucleotide[i] for i in range(len(true_seq))]
+#     acceptor_probs = [splicing_pred[i] * acceptor_dinucleotide[i] for i in range(len(true_seq))]
+#     return donor_probs, acceptor_probs
+#

geney/utils/spliceai_utils.py

@@ -0,0 +1,123 @@
+# __all__ = ['sai_predict_probs']
+#
+# #### SpliceAI Modules
+#
+# from keras.models import load_model
+# from importlib import resources
+# import numpy as np
+# import tensorflow as tf
+# import sys
+# import absl.logging
+# absl.logging.set_verbosity(absl.logging.ERROR)
+#
+#
+# # Check if GPU is available
+# # if tf.config.list_physical_devices('GPU'):
+# #     print("Running on GPU.")
+# # else:
+# #     print("Running on CPU.")
+#
+# # tf.config.threading.set_intra_op_parallelism_threads(1)
+# # tf.config.threading.set_inter_op_parallelism_threads(1)
+#
+# # List the model filenames relative to the spliceai package.
+# model_filenames = [f"models/spliceai{i}.h5" for i in range(1, 6)]
+#
+# # Load each model using the package resources.
+# # sai_models = [load_model(resources.files("spliceai").joinpath(filename))
+# #               for filename in model_filenames]
+#
+# if sys.platform == 'darwin':
+#     sai_paths = ('models/spliceai{}.h5'.format(x) for x in range(1, 6))
+#     # sai_models = [load_model(resource_filename('spliceai', x)) for x in sai_paths]
+#     sai_models = [load_model(resources.files('spliceai').joinpath(f)) for f in sai_paths]
+# else:
+#     sai_paths = ['/tamir2/nicolaslynn/home/miniconda3/lib/python3.10/site-packages/spliceai/models/spliceai1.h5',
+#                  '/tamir2/nicolaslynn/home/miniconda3/lib/python3.10/site-packages/spliceai/models/spliceai2.h5',
+#                  '/tamir2/nicolaslynn/home/miniconda3/lib/python3.10/site-packages/spliceai/models/spliceai3.h5',
+#                  '/tamir2/nicolaslynn/home/miniconda3/lib/python3.10/site-packages/spliceai/models/spliceai4.h5',
+#                  '/tamir2/nicolaslynn/home/miniconda3/lib/python3.10/site-packages/spliceai/models/spliceai5.h5']
+#
+#     sai_models = [load_model(f) for f in sai_paths]
+
+__all__ = ['sai_predict_probs']
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+
+import absl.logging
+absl.logging.set_verbosity(absl.logging.ERROR)
+
+import os
+import sys
+import tensorflow as tf
+import numpy as np
+from keras.models import load_model
+from importlib import resources
+
+
+# Force device selection
+if tf.config.list_physical_devices('GPU'):
+    device = '/GPU:0'
+elif sys.platform == 'darwin' and tf.config.list_physical_devices('MPS'):
+    device = '/device:GPU:0'  # MPS uses /device:GPU:0 in TF
+else:
+    device = '/CPU:0'
+
+# Model loading paths
+if sys.platform == 'darwin':
+    model_filenames = [f"models/spliceai{i}.h5" for i in range(1, 6)]
+    model_paths = [resources.files('spliceai').joinpath(f) for f in model_filenames]
+else:
+    model_paths = [f"/tamir2/nicolaslynn/home/miniconda3/lib/python3.10/site-packages/spliceai/models/spliceai{i}.h5"
+                   for i in range(1, 6)]
+
+# Load models onto correct device
+with tf.device(device):
+    sai_models = [load_model(str(f)) for f in model_paths]
+
+
+print(f"SpliceAI loaded to {device}.")
+
+def one_hot_encode(seq):
+
+    map = np.asarray([[0, 0, 0, 0],
+                      [1, 0, 0, 0],
+                      [0, 1, 0, 0],
+                      [0, 0, 1, 0],
+                      [0, 0, 0, 1]])
+
+    seq = seq.upper().replace('A', '\x01').replace('C', '\x02')
+    seq = seq.replace('G', '\x03').replace('T', '\x04').replace('N', '\x00')
+
+    return map[np.fromstring(seq, np.int8) % 5]
+
+
+def sai_predict_probs(seq: str, models: list) -> list:
+    '''
+    Predicts the donor and acceptor junction probability of each
+    NT in seq using SpliceAI.
+
+    Let m:=2*sai_mrg_context + L be the input seq length. It is assumed
+    that the input seq has the following structure:
+
+          seq = |<sai_mrg_context NTs><L NTs><sai_mrg_context NTs>|
+
+    The returned probability matrix is of size 2XL, where
+    the first row is the acceptor probability and the second row
+    is the donor probability. These probabilities corresponds to the
+    middel <L NTs> NTs of the input seq.
+    '''
+    x = one_hot_encode(seq)[None, :]
+    y = np.mean([models[m].predict(x, verbose=0) for m in range(5)], axis=0)
+    # return y[0, :, 1:].T
+    y = y[0, :, 1:].T
+    return y[0, :], y[1, :]
+
+
+def run_spliceai_seq(seq, indices, threshold=0):
+    # seq = 'N' * 5000 + seq + 'N' * 5000
+    ref_seq_probs_temp = sai_predict_probs(seq, sai_models)
+    ref_seq_acceptor_probs, ref_seq_donor_probs = ref_seq_probs_temp[0, :], ref_seq_probs_temp[1, :]
+    acceptor_indices = {a: b for a, b in list(zip(indices, ref_seq_acceptor_probs)) if b >= threshold}
+    donor_indices = {a: b for a, b in list(zip(indices, ref_seq_donor_probs)) if b >= threshold}
+    return donor_indices, acceptor_indices