[petsc-users] Orthogonalization of a (sparse) PETSc matrix

[Thanasis Boutsikakis]

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,
)
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://lists.mcs.anl.gov/pipermail/petsc-users/attachments/20230829/98d3c874/attachment-0001.html>