Array data dependence analysis

loki/analyse/util_linear_algebra.py

@@ -5,9 +5,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-from numpy import zeros_like, dot
-from numpy import vstack, hstack
-from numpy import all as np_all, sum as np_sum, isin as np_isin
+import numpy as np
 
 __all__ = [
     "back_substitution",
@@ -21,75 +19,78 @@ def is_independent_system(matrix):
     """
     Check if a linear system of equations can be split into independent one-dimensional problems.
 
-    Args:
-        matrix (numpy.ndarray): A rectangular matrix representing coefficients.
+    Parameters
+    ----------
+    matrix : numpy.ndarray
+        A rectangular matrix representing coefficients.
 
-    Returns:
-        bool: True if the system can be split into independent one-dimensional problems, False otherwise.
+    Returns
+    -------
+    bool
+        True if the system can be split into independent one-dimensional problems, False otherwise.
 
-    This function checks whether a linear system of equations in the form of matrix x [operator] right_hand_side
+    Notes
+    -----
+    This function checks whether a linear system of equations in the form of `matrix [operator] right_hand_side`
     can be split into independent one-dimensional problems. The number of problems is determined by the
     number of variables (the row number of the matrix).
 
     Each problem consists of a coefficient vector and a right-hand side. The system can be considered independent
     if each row of the matrix has exactly one non-zero coefficient or no non-zero coefficients.
     """
 
-    return np_all(np_isin(np_sum(matrix != 0, axis=1), [0, 1]))
+    return np.all(np.isin(np.sum(matrix != 0, axis=1), [0, 1]))
 
 
 def yield_one_d_systems(matrix, right_hand_side):
     """
-    Split a linear system of equations (<=, >= or ==) into independent one-dimensional problems.
-
-    Args:
-        matrix (numpy.ndarray): A rectangular matrix representing coefficients.
-        right_hand_side (numpy.ndarray): The right-hand side vector.
-
-    Yields:
-        tuple[single_dimensional_array, single_dimensional_array]:
-            A tuple containing a coefficient vector and the corresponding right-hand side.
-
-    This function takes a linear system of equations in the form of matrix x [operator] right_hand_side,
-    where "matrix" is a rectangular matrix, x is a vector of variables, and "right_hand_side" is
+    Split a linear system of equations (<=, >=, or ==) into independent one-dimensional problems.
+
+    Parameters
+    ----------
+    matrix : numpy.ndarray
+        A rectangular matrix representing coefficients.
+    right_hand_side : numpy.ndarray
+        The right-hand side vector.
+
+    Yields
+    ------
+    tuple[numpy.ndarray, numpy.ndarray]
+        A tuple containing a coefficient vector and the corresponding right-hand side.
+
+    Notes
+    -----
+    The independence of the problems is NOT explicitly checked; call `is_independent_system` before using this
+    function if unsure.
+
+    This function takes a linear system of equations in the form of `matrix [operator] right_hand_side`,
+    where "matrix" is a rectangular matrix, "x" is a vector of variables, and "right_hand_side" is
     the right-hand side vector. It splits the system into assumed independent one-dimensional problems.
 
     Each problem consists of a coefficient vector and a right-hand side. The number of problems is equal to the
     number of variables (the row number of the matrix).
 
-    Note:
-    - The independence of the problems is not explicitly checked, call is_independent_system before using this
-      function if unsure.
-
-    Example:
-    ```
+    Example
+    -------
+    ```python
     for A, b in yield_one_d_systems(matrix, right_hand_side):
         # Solve the one-dimensional problem A * x = b
         solution = solve_one_d_system(A, b)
     ```
     """
-    # drop completly empty rows
-    mask = ~np_all(hstack((matrix, right_hand_side)) == 0, axis=1)
-    matrix = matrix[mask]
-    right_hand_side = right_hand_side[mask]
-
     # yield systems with empty left hand side (A) and non empty right hand side
-    mask = np_all(matrix == 0, axis=1)
-    if right_hand_side[mask].size == 0:
-        return
-
-    for A, b in zip(matrix[mask].T, right_hand_side[mask].T):
-        yield A, b
+    mask = np.all(matrix == 0, axis=1)
+    if right_hand_side[mask].size != 0:
+        for A in matrix[mask].T:
+            yield A.reshape((-1,1)), right_hand_side[mask]
 
     matrix = matrix[~mask]
     right_hand_side = right_hand_side[~mask]
 
-    if right_hand_side.size == 0:
-        return
-
-    for A, b in zip(matrix.T, right_hand_side.T):
-        mask = A != 0
-        yield A[mask], b[mask]
+    if right_hand_side.size != 0:
+        for A in matrix.T:
+            mask = A != 0
+            yield A[mask].reshape((-1,1)), right_hand_side[mask]
 
 
 def back_substitution(
@@ -100,33 +101,42 @@ def back_substitution(
     """
     Solve a linear system of equations using back substitution for an upper triangular square matrix.
 
-    Args:
-        upper_triangular_square_matrix (numpy.ndarray): An upper triangular square matrix (R).
-        right_hand_side (numpy.ndarray): A vector (y) on the right-hand side of the equation Rx = y.
-        division_operation (function, optional): A custom division operation function. Default is standard division (/).
+    Parameters
+    ----------
+    upper_triangular_square_matrix : numpy.ndarray
+        An upper triangular square matrix (R).
+
+    right_hand_side : numpy.ndarray
+        A vector (y) on the right-hand side of the equation Rx = y.
 
-    Returns:
-        numpy.ndarray: The solution vector (x) to the system of equations Rx = y.
+    division_operation : function, optional
+        A custom division operation function. Default is standard division (/).
 
+    Returns
+    -------
+    numpy.ndarray
+        The solution vector (x) to the system of equations Rx = y.
+
+    Notes
+    -----
     The function performs back substitution to find the solution vector x for the equation Rx = y,
     where R is an upper triangular square matrix and y is a vector. The division_operation
     function is used for division (e.g., for custom division operations).
 
-    Note:
-    - The function assumes that the upper right element of the upper_triangular_square_matrix (R)
-      is nonzero for proper back substitution.
+    The function assumes that the upper right element of the upper_triangular_square_matrix (R)
+    is nonzero for proper back substitution.
     """
     R = upper_triangular_square_matrix
     y = right_hand_side
 
-    x = zeros_like(y)
+    x = np.zeros_like(y)
 
     assert R[-1, -1] != 0
 
     x[-1] = divison_operation(y[-1], R[-1, -1])
 
     for i in range(len(y) - 2, -1, -1):
-        x[i] = divison_operation((y[i] - dot(R[i, i + 1 :], x[i + 1 :])), R[i, i])
+        x[i] = divison_operation((y[i] - np.dot(R[i, i + 1 :], x[i + 1 :])), R[i, i])
 
     return x
 
@@ -135,25 +145,32 @@ def generate_row_echelon_form(
     A, conditional_check=lambda A: None, division_operator=lambda x, y: x / y
 ):
     """
-    Calculate the Row Echelon Form (REF) of a matrix A using Gaussian elimination.
-
-    Args:
-        A (numpy.ndarray): The input matrix for which RREF is to be calculated.
-        conditional_check (function, optional): A custom function to check conditions during the computation.
-        division_operation (function, optional): A custom division operation function. Default is standard division (/).
-
-    Returns:
-        numpy.ndarray: The REF of the input matrix A.
-
-    This function computes the Row Echelon Form (RREF) of a given matrix A using Gaussian elimination.
-    You can provide a custom conditional_check function to add checks or operations during the computation.
-    You can also specify a custom division_operation function for division (e.g., for custom division operations).
-
-    Note:
+    Calculate the Row Echelon Form (REF) of a matrix.
+
+    Parameters
+    ----------
+    A : numpy.ndarray
+        The input matrix for which the REF is to be calculated.
+    conditional_check : function, optional
+        A custom function to check conditions during the computation.
+    division_operation : function, optional
+        A custom division operation function. Default is standard division (/).
+
+    Returns
+    -------
+    numpy.ndarray
+        The REF of the input matrix A.
+
+    Notes
+    -----
     - If the input matrix has no rows or columns, it is already in REF, and the function returns itself.
     - The function utilizes the specified division operation (default is standard division) for division.
 
-    Reference: https://math.stackexchange.com/questions/3073083/how-to-reduce-matrix-into-row-echelon-form-in-numpy
+    Reference
+    ---------
+    https://math.stackexchange.com/a/3073117
+    for question:
+    https://math.stackexchange.com/questions/3073083/how-to-reduce-matrix-into-row-echelon-form-in-numpy
     """
     # if matrix A has no columns or rows,
     # it is already in REF, so we return itself
@@ -170,7 +187,7 @@ def generate_row_echelon_form(
         # we perform REF on matrix from second column
         B = generate_row_echelon_form(A[:, 1:], conditional_check, division_operator)
         # and then add the first zero-column back
-        return hstack([A[:, :1], B])
+        return np.hstack([A[:, :1], B])
 
     # if non-zero element happens not in the first row,
     # we switch rows
@@ -190,4 +207,4 @@ def generate_row_echelon_form(
     B = generate_row_echelon_form(A[1:, 1:], conditional_check, division_operator)
 
     # we add first row and first (zero) column, and return
-    return vstack([A[:1], hstack([A[1:, :1], B])])
+    return np.vstack([A[:1], np.hstack([A[1:, :1], B])])