njit added back

caretta/dynamic_time_warping.py

@@ -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.):
     """

caretta/helper.py

@@ -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])

caretta/msa_numba.py

@@ -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)

caretta/neighbor_joining.py

@@ -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):
     """

caretta/psa_numba.py

@@ -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,

caretta/rmsd_calculations.py

@@ -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:
     """