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,45 @@ def jacobi_iteration_method(
             temp = (temp + val) / denom
             new_val.append(temp)
         init_val = new_val
+    """
+
+    """
+    denom - a list of values along the diagonal
+    """
+    denom = 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_diag - coefficient_matrix array values without diagonal elements
+    """
+    no_diag = 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)
+        temp = np.sum((-1) * no_diag * arr, axis=1)
+        new_val = (temp + val_last) / denom
+        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