jacobi_iteration_method.py the use of vector operations, which reduces the calculation time by dozens of times

arithmetic_analysis/jacobi_iteration_method.py

@@ -12,7 +12,7 @@
 def jacobi_iteration_method(
     coefficient_matrix: NDArray[float64],
     constant_matrix: NDArray[float64],
-    init_val: list[int],
+    init_val: list[float],
     iterations: int,
 ) -> list[float]:
     """
@@ -115,6 +115,7 @@ def jacobi_iteration_method(
 
     strictly_diagonally_dominant(table)
 
+    """
     # Iterates the whole matrix for given number of times
     for _ in range(iterations):
         new_val = []
@@ -130,8 +131,37 @@ def jacobi_iteration_method(
             temp = (temp + val) / denom
             new_val.append(temp)
         init_val = new_val
+    """
+
+    # denominator - a list of values along the diagonal
+    denominator = np.diag(coefficient_matrix)
+
+    # val_last - values of the last column of the table array
+    val_last = table[:, -1]
+
+    # masks - boolean mask of all strings without diagonal
+    # elements array coefficient_matrix
+    masks = ~np.eye(coefficient_matrix.shape[0], dtype=bool)
+
+    # no_diagonals - coefficient_matrix array values without diagonal elements
+    no_diagonals = coefficient_matrix[masks].reshape(-1, rows - 1)
+
+    # Here we get 'i_col' - these are the column numbers, for each row
+    # without diagonal elements, except for the last column.
+    i_row, i_col = np.where(masks)
+    ind = i_col.reshape(-1, rows - 1)
+
+    #'i_col' is converted to a two-dimensional list 'ind', which will be
+    # used to make selections from 'init_val' ('arr' array see below).
+
+    # Iterates the whole matrix for given number of times
+    for _ in range(iterations):
+        arr = np.take(init_val, ind)
+        sum_product_rows = np.sum((-1) * no_diagonals * arr, axis=1)
+        new_val = (sum_product_rows + val_last) / denominator
+        init_val = new_val
 
-    return [float(i) for i in new_val]
+    return new_val.tolist()
 
 
 # Checks if the given matrix is strictly diagonally dominant