Hi all, I have the following code that orthogonalizes a PETSc matrix. The problem is that this implementation requires that the PETSc matrix is dense, otherwise, it fails at bv.SetFromOptions(). Hence the assert in orthogonality().
What could I do in order to be able to orthogonalize sparse matrices as well? Could I convert it efficiently? (I tried to no avail) Thanks! """Experimenting with matrix orthogonalization""" import contextlib import sys import time import numpy as np from firedrake import COMM_WORLD from firedrake.petsc import PETSc import slepc4py slepc4py.init(sys.argv) from slepc4py import SLEPc from numpy.testing import assert_array_almost_equal EPSILON_USER = 1e-4 EPS = sys.float_info.epsilon def Print(message: str): """Print function that prints only on rank 0 with color Args: message (str): message to be printed """ PETSc.Sys.Print(message) def create_petsc_matrix(input_array, sparse=True): """Create a PETSc matrix from an input_array Args: input_array (np array): Input array partition_like (PETSc mat, optional): Petsc matrix. Defaults to None. sparse (bool, optional): Toggle for sparese or dense. Defaults to True. Returns: PETSc mat: PETSc matrix """ # Check if input_array is 1D and reshape if necessary assert len(input_array.shape) == 2, "Input array should be 2-dimensional" global_rows, global_cols = input_array.shape size = ((None, global_rows), (global_cols, global_cols)) # Create a sparse or dense matrix based on the 'sparse' argument if sparse: matrix = PETSc.Mat().createAIJ(size=size, comm=COMM_WORLD) else: matrix = PETSc.Mat().createDense(size=size, comm=COMM_WORLD) matrix.setUp() local_rows_start, local_rows_end = matrix.getOwnershipRange() for counter, i in enumerate(range(local_rows_start, local_rows_end)): # Calculate the correct row in the array for the current process row_in_array = counter + local_rows_start matrix.setValues( i, range(global_cols), input_array[row_in_array, :], addv=False ) # Assembly the matrix to compute the final structure matrix.assemblyBegin() matrix.assemblyEnd() return matrix def orthogonality(A): # sourcery skip: avoid-builtin-shadow """Checking and correcting orthogonality Args: A (PETSc.Mat): Matrix of size [m x k]. Returns: PETSc.Mat: Matrix of size [m x k]. """ # Check if the matrix is dense mat_type = A.getType() assert mat_type in ( "seqdense", "mpidense", ), "A must be a dense matrix. SLEPc.BV().createFromMat() requires a dense matrix." m, k = A.getSize() Phi1 = A.getColumnVector(0) Phi2 = A.getColumnVector(k - 1) # Compute dot product using PETSc function dot_product = Phi1.dot(Phi2) if abs(dot_product) > min(EPSILON_USER, EPS * m): Print(" Matrix is not orthogonal") # Type can be CHOL, GS, mro(), SVQB, TSQR, TSQRCHOL _type = SLEPc.BV().OrthogBlockType.GS bv = SLEPc.BV().createFromMat(A) bv.setFromOptions() bv.setOrthogonalization(_type) bv.orthogonalize() A = bv.createMat() Print(" Matrix successfully orthogonalized") # # Assembly the matrix to compute the final structure if not A.assembled: A.assemblyBegin() A.assemblyEnd() else: Print(" Matrix is orthogonal") return A # -------------------------------------------- # EXP: Orthogonalization of an mpi PETSc matrix # -------------------------------------------- m, k = 11, 7 # Generate the random numpy matrices np.random.seed(0) # sets the seed to 0 A_np = np.random.randint(low=0, high=6, size=(m, k)) A = create_petsc_matrix(A_np, sparse=False) A_orthogonal = orthogonality(A) # -------------------------------------------- # TEST: Orthogonalization of a numpy matrix # -------------------------------------------- # Generate A_np_orthogonal A_np_orthogonal, _ = np.linalg.qr(A_np) # Get the local values from A_orthogonal local_rows_start, local_rows_end = A_orthogonal.getOwnershipRange() A_orthogonal_local = A_orthogonal.getValues( range(local_rows_start, local_rows_end), range(k) ) # Assert the correctness of the multiplication for the local subset assert_array_almost_equal( np.abs(A_orthogonal_local), np.abs(A_np_orthogonal[local_rows_start:local_rows_end, :]), decimal=5, )