Commit b63dd2ef authored by Akdel's avatar Akdel
Browse files

njit added back

parent 2774b404
......@@ -131,7 +131,7 @@ def _get_dtw_alignment(start_direction, backtrack: np.ndarray, n1, m1):
return indices_1[:index][::-1], indices_2[:index][::-1]
# @nb.njit
@nb.njit
# @numba_cc.export('dtw_align', '(f64[:], f64, f64)')
def dtw_align(distance_matrix: np.ndarray, gap_open_penalty: float = 0., gap_extend_penalty: float = 0.):
"""
......
......@@ -8,7 +8,7 @@ import numpy as np
import prody as pd
# @nb.njit
@nb.njit
# @numba_cc.export('nan_normalize', 'f64[:](f64[:])')
def nan_normalize(numbers):
minv, maxv = np.nanmin(numbers), np.nanmax(numbers)
......@@ -42,7 +42,7 @@ def aligned_string_to_array(aln: str) -> np.ndarray:
return aln_array
# @nb.njit
@nb.njit
# @numba_cc.export('get_common_positions', '(i64[:], i64[:], i64)')
def get_common_positions(aln_array_1, aln_array_2, gap=-1):
"""
......@@ -63,7 +63,7 @@ def get_common_positions(aln_array_1, aln_array_2, gap=-1):
return pos_1, pos_2
# @nb.njit
@nb.njit
# @numba_cc.export('get_aligned_data', '(i64[:], f64[:], i64)')
def get_aligned_data(aln_array: np.ndarray, data: np.ndarray, gap=-1):
"""
......@@ -90,7 +90,7 @@ def get_aligned_data(aln_array: np.ndarray, data: np.ndarray, gap=-1):
return aln_coords
# @nb.njit
@nb.njit
# @numba_cc.export('get_aligned_string_data', '(i64[:], i8[:], i64)')
def get_aligned_string_data(aln_array, data, gap=-1):
pos = np.array([i for i in range(len(aln_array)) if aln_array[i] != gap])
......
......@@ -13,14 +13,14 @@ from caretta import psa_numba as psa
from caretta import rmsd_calculations, helper
# @nb.njit
@nb.njit
def get_common_coordinates(coords_1, coords_2, aln_1, aln_2, gap=-1):
assert aln_1.shape == aln_2.shape
pos_1, pos_2 = helper.get_common_positions(aln_1, aln_2, gap)
return coords_1[pos_1], coords_2[pos_2]
# @nb.njit
@nb.njit
def make_pairwise_dtw_score_matrix(coords_array, secondary_array, lengths_array, gamma,
gap_open_penalty: float, gap_extend_penalty: float,
gap_open_sec, gap_extend_sec):
......@@ -45,16 +45,13 @@ def make_pairwise_dtw_score_matrix(coords_array, secondary_array, lengths_array,
return pairwise_matrix
# @nb.njit(parallel=True)
@nb.njit(parallel=True)
def make_pairwise_rmsd_score_matrix(coords_array, secondary_array, lengths_array, gamma,
gap_open_penalty: float, gap_extend_penalty: float,
gap_open_sec, gap_extend_sec):
pairwise_matrix = np.zeros((coords_array.shape[0], coords_array.shape[0]))
for i in nb.prange(pairwise_matrix.shape[0] - 1):
for j in range(i + 1, pairwise_matrix.shape[1]):
print(i, j)
print(coords_array[i, :lengths_array[i]], coords_array[j, :lengths_array[j]],
secondary_array[i, :lengths_array[i]], secondary_array[j, :lengths_array[j]])
dtw_aln_1, dtw_aln_2, score = psa.get_pairwise_alignment(coords_array[i, :lengths_array[i]], coords_array[j, :lengths_array[j]],
secondary_array[i, :lengths_array[i]], secondary_array[j, :lengths_array[j]],
gamma,
......@@ -73,7 +70,7 @@ def make_pairwise_rmsd_score_matrix(coords_array, secondary_array, lengths_array
return pairwise_matrix
# @nb.njit
@nb.njit
def _get_alignment_data(coords_1, coords_2, secondary_1, secondary_2, gamma,
gap_open_sec, gap_extend_sec,
gap_open_penalty: float, gap_extend_penalty: float):
......@@ -95,7 +92,7 @@ def _get_alignment_data(coords_1, coords_2, secondary_1, secondary_2, gamma,
return aln_coords_1, aln_coords_2, aln_sec_1, aln_sec_2, dtw_aln_1, dtw_aln_2
# @nb.njit
@nb.njit
def get_mean_coords_extra(aln_coords_1: np.ndarray, aln_coords_2: np.ndarray) -> np.ndarray:
"""
Mean of two coordinate sets (of the same shape)
......@@ -119,7 +116,7 @@ def get_mean_coords_extra(aln_coords_1: np.ndarray, aln_coords_2: np.ndarray) ->
return mean_coords
# @nb.njit
@nb.njit
def get_mean_secondary(aln_sec_1: np.ndarray, aln_sec_2: np.ndarray, gap=0) -> np.ndarray:
"""
Mean of two coordinate sets (of the same shape)
......
......@@ -13,7 +13,7 @@ import numpy as np
# 5. Start the algorithm again, replacing the pair of joined neighbors with the new node and using the distances calculated in the previous step.
# @nb.njit
@nb.njit
# @numba_cc.export('neighbor_joining', '(f64[:])')
def neighbor_joining(distance_matrix: np.ndarray) -> (np.ndarray, np.ndarray):
"""
......@@ -88,7 +88,7 @@ def neighbor_joining(distance_matrix: np.ndarray) -> (np.ndarray, np.ndarray):
# Q matrix calculation + minimum i, j (step 1 & 2)
# @nb.njit
@nb.njit
# @numba_cc.export('_find_join_nodes', '(f64[:])')
def _find_join_nodes(distance_matrix):
"""
......@@ -119,7 +119,7 @@ def _find_join_nodes(distance_matrix):
# Branch length calculation (step 3)
# @nb.njit
@nb.njit
# @numba_cc.export('_find_branch_length', '(f64[:], i64, i64)')
def _find_branch_length(distance_matrix, i, j):
"""
......
......@@ -5,7 +5,7 @@ from caretta import dynamic_time_warping as dtw
from caretta import rmsd_calculations, helper
# @nb.njit
@nb.njit
# @numba_cc.export('make_signal_index', 'f64[:](f64[:], i64)')
def make_signal_index(coords, index):
centroid = coords[index]
......@@ -15,7 +15,7 @@ def make_signal_index(coords, index):
return distances
# @nb.njit
@nb.njit
# @numba_cc.export('dtw_signals_index', '(f64[:], f64[:], i64, i64, i64)')
def dtw_signals_index(coords_1, coords_2, index, size=30, overlap=1):
signals_1 = np.zeros(((coords_1.shape[0] - size) // overlap, size))
......@@ -46,7 +46,7 @@ def dtw_signals_index(coords_1, coords_2, index, size=30, overlap=1):
return coords_1, coords_2
# @nb.njit
@nb.njit
# @numba_cc.export('get_dtw_signal_score_pos', '(f64[:], f64[:], f64, i64, f64, f64)')
def get_dtw_signal_score_pos(coords_1, coords_2, gamma, index, gap_open_penalty, gap_extend_penalty):
coords_1[:, :3], coords_2[:, :3] = dtw_signals_index(coords_1[:, :3], coords_2[:, :3], index)
......@@ -61,7 +61,7 @@ def get_dtw_signal_score_pos(coords_1, coords_2, gamma, index, gap_open_penalty,
return rmsd_calculations.get_caretta_score(common_coords_1, common_coords_2, gamma, False), pos_1, pos_2
# @nb.njit
@nb.njit
# @numba_cc.export('get_secondary_distance_matrix', '(i8[:], i8[:], i8)')
def get_secondary_distance_matrix(secondary_1, secondary_2, gap=0):
score_matrix = np.zeros((secondary_1.shape[0], secondary_2.shape[0]))
......@@ -75,7 +75,7 @@ def get_secondary_distance_matrix(secondary_1, secondary_2, gap=0):
return score_matrix
# @nb.njit
@nb.njit
# @numba_cc.export('get_secondary_rmsd_pos', '(i8[:], i8[:], f64[:], f64[:], f64, f64, f64)')
def get_secondary_rmsd_pos(secondary_1, secondary_2, coords_1, coords_2, gamma, gap_open_sec, gap_extend_sec):
distance_matrix = get_secondary_distance_matrix(secondary_1, secondary_2)
......@@ -87,7 +87,7 @@ def get_secondary_rmsd_pos(secondary_1, secondary_2, coords_1, coords_2, gamma,
return rmsd_calculations.get_caretta_score(common_coords_1, common_coords_2, gamma, False), pos_1, pos_2
# @nb.njit
@nb.njit
# @numba_cc.export('get_pairwise_alignment', '(f64[:], f64[:], i8[:], i8[:], f64, f64, f64, f64, f64, i64)')
def get_pairwise_alignment(coords_1, coords_2,
secondary_1, secondary_2,
......
......@@ -2,14 +2,14 @@ import numba as nb
import numpy as np
# @nb.njit
@nb.njit
# @numba_cc.export('normalize', 'f64[:](f64[:])')
def normalize(numbers):
minv, maxv = np.min(numbers), np.max(numbers)
return (numbers - minv) / (maxv - minv)
# @nb.njit
@nb.njit
# @numba_cc.export('nb_mean_axis_0', 'f64[:](f64[:])')
def nb_mean_axis_0(array: np.ndarray) -> np.ndarray:
"""
......@@ -21,7 +21,7 @@ def nb_mean_axis_0(array: np.ndarray) -> np.ndarray:
return mean_array
# @nb.njit
@nb.njit
# @numba_cc.export('svd_superimpose', '(f64[:], f64[:])')
def svd_superimpose(coords_1: np.ndarray, coords_2: np.ndarray):
"""
......@@ -51,7 +51,7 @@ def svd_superimpose(coords_1: np.ndarray, coords_2: np.ndarray):
return rotation_matrix.astype(np.float64), translation_matrix.astype(np.float64)
# @nb.njit
@nb.njit
# @numba_cc.export('apply_rotran', '(f64[:], f64[:], f64[:])')
def apply_rotran(coords: np.ndarray, rotation_matrix: np.ndarray, translation_matrix: np.ndarray) -> np.ndarray:
"""
......@@ -71,7 +71,7 @@ def apply_rotran(coords: np.ndarray, rotation_matrix: np.ndarray, translation_ma
# @numba_cc.export('superpose_with_pos', '(f64[:], f64[:], f64[:], f64[:])')
# @nb.njit
@nb.njit
def superpose_with_pos(coords_1, coords_2, common_coords_1, common_coords_2):
"""
Superpose two sets of un-aligned coordinates using smaller subsets of aligned coordinates
......@@ -94,7 +94,7 @@ def superpose_with_pos(coords_1, coords_2, common_coords_1, common_coords_2):
return coords_1, coords_2, common_coords_2_rot
# @nb.njit
@nb.njit
# @numba_cc.export('make_distance_matrix', '(f64[:], f64[:], f64, b1)')
def make_distance_matrix(coords_1: np.ndarray, coords_2: np.ndarray, gamma, normalized=False) -> np.ndarray:
"""
......@@ -122,7 +122,7 @@ def make_distance_matrix(coords_1: np.ndarray, coords_2: np.ndarray, gamma, norm
return distance_matrix
# @nb.njit
@nb.njit
# @numba_cc.export('get_rmsd', '(f64[:], f64[:])')
def get_rmsd(coords_1: np.ndarray, coords_2: np.ndarray) -> float:
"""
......@@ -131,7 +131,7 @@ def get_rmsd(coords_1: np.ndarray, coords_2: np.ndarray) -> float:
return np.sqrt(np.sum((coords_1 - coords_2) ** 2) / coords_1.shape[0])
# @nb.njit
@nb.njit
# @numba_cc.export('get_caretta_score', '(f64[:], f64[:], f64, b1)')
def get_caretta_score(coords_1: np.ndarray, coords_2: np.ndarray, gamma, normalized) -> float:
"""
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment