added mean function for shapes

caretta/multiple_alignment.py

@@ -22,8 +22,7 @@ def get_common_coordinates(coords_1, coords_2, aln_1, aln_2, gap=-1):
 
 
 @nb.njit
-def get_mean_coords(aln_coords_1: np.ndarray, aln_coords_2: np.ndarray) -> typing.Tuple[
-    np.ndarray, np.ndarray]:
+def get_mean_coords(aln_coords_1: np.ndarray, aln_coords_2: np.ndarray) -> np.ndarray:
     """
     Mean of two coordinate sets (of the same shape)
     Last column has the consensus weight * number of aligned residues
@@ -57,9 +56,11 @@ def get_mean_weights(aln_coords_1: np.ndarray, aln_coords_2: np.ndarray, aln_wei
 
 def get_mean_shape_array(shape_array1: np.ndarray, shape_array2: np.ndarray, n_bits=64, take_first=False):
     mean_shape = np.zeros(shape_array1.shape[0], dtype=shape_array1.dtype)
-    for i in range(shape_array1.shape[0]):
-        s1 = shape_array1.view(f"int{n_bits}")[i]
-        s2 = shape_array2[i]
+    shape_array1_flat = shape_array1.flatten()
+    shape_array2_flat = shape_array2.flatten()
+    for i in range(shape_array1_flat.shape[0]):
+        s1 = shape_array1_flat.view(f"int{n_bits}")[i]
+        s2 = shape_array2_flat.view(f"int{n_bits}")[i]
         if s1 == -1:
             mean_shape[i] = s2
         elif s2 == -1:
@@ -69,7 +70,7 @@ def get_mean_shape_array(shape_array1: np.ndarray, shape_array2: np.ndarray, n_b
                 mean_shape[i] = s1
             else:
                 mean_shape[i] = mean_two_shapes(s1, s2, n_bits=n_bits)
-    return mean_shape
+    return mean_shape.reshape((-1, 1))
 
 
 def mean_two_shapes(s1, s2, n_bits=64):