PETScMatrix.h

Go to the documentation of this file.

// SPDX-FileCopyrightText: Copyright (c) OpenGeoSys Community (opengeosys.org)
// SPDX-License-Identifier: BSD-3-Clause
 
#pragma once
 
#include <petscmat.h>
 
#include <string>
#include <vector>
 
#include "MathLib/LinAlg/RowColumnIndices.h"
#include "MathLib/LinAlg/SparsityPattern.h"
 
typedef Mat PETSc_Mat;
 
namespace MathLib
{
class PETScVector;
class PETScMatrix
{
public:
    using IndexType = PetscInt;
 
public:
    PETScMatrix() {}
    PETScMatrix(const PetscInt nrows,
                const PETScSparsityPattern& sparsity_pattern);

    PETScMatrix(const PetscInt nrows, const PetscInt ncols,
                const PETScSparsityPattern& sparsity_pattern);
 
    ~PETScMatrix() { destroy(); }
    PETScMatrix(PETScMatrix const& A);
 
    PETScMatrix& operator=(PETScMatrix const& A);

    void finalizeAssembly(const MatAssemblyType asm_type = MAT_FINAL_ASSEMBLY)
    {
        MatAssemblyBegin(A_, asm_type);
        MatAssemblyEnd(A_, asm_type);
    }

    PetscInt getNumberOfRows() const { return nrows_; }
    PetscInt getNumberOfColumns() const { return ncols_; }
    PetscInt getNumberOfLocalRows() const { return n_loc_rows_; }
    PetscInt getNumberOfLocalColumns() const { return n_loc_cols_; }
    PetscInt getRangeBegin() const { return start_rank_; }
    PetscInt getRangeEnd() const { return end_rank_; }
    std::vector<PetscInt> const& getSparsityColumnIndices() const
    {
        return sparsity_col_idx_;
    }

    Mat& getRawMatrix() { return A_; }
    Mat const& getRawMatrix() const { return A_; }
    void setZero() { MatZeroEntries(A_); }
    void setRowsColumnsZero(std::vector<PetscInt> const& row_pos);

    void set(const PetscInt i, const PetscInt j, const PetscScalar value)
    {
        PetscCallAbort(PETSC_COMM_WORLD,
                       MatSetValue(A_, i, j, value, INSERT_VALUES));
    }

    void add(const PetscInt i, const PetscInt j, const PetscScalar value)
    {
        PetscCallAbort(PETSC_COMM_WORLD,
                       MatSetValue(A_, i, j, value, ADD_VALUES));
    }

    void addToDiagonal(const PetscScalar value);

    template <class T_DENSE_MATRIX>
    void add(RowColumnIndices<PetscInt> const& indices,
             const T_DENSE_MATRIX& sub_matrix)
    {
        // Set global column indices to positive to allow all entries of columns
        // to be added to the global matrix. For the ghost columns, only the
        // off diagonal entries are added due to the negative indices of the
        // corresponding rows.
        std::vector<PetscInt> cols;
        cols.reserve(indices.columns.size());
        for (auto col : indices.columns)
        {
            // Ghost entries, and its original index is 0.
            if (col == -ncols_)
            {
                cols.push_back(0);
            }
            else
            {
                cols.push_back(std::abs(col));
            }
        }
 
        add(indices.rows, cols, sub_matrix);
    }

    template <class T_DENSE_MATRIX>
    void add(std::vector<PetscInt> const& row_pos,
             std::vector<PetscInt> const& col_pos,
             const T_DENSE_MATRIX& sub_mat);

    void viewer(const std::string& file_name,
                const PetscViewerFormat vw_format = PETSC_VIEWER_ASCII_MATLAB);
 
private:
    void destroy()
    {
        if (A_ != nullptr)
        {
            PetscCallAbort(PETSC_COMM_WORLD, MatDestroy(&A_));
        }
        A_ = nullptr;
    }

    Mat A_ = nullptr;

    PetscInt nrows_;

    PetscInt ncols_;

    PetscInt n_loc_rows_;

    PetscInt n_loc_cols_;

    PetscInt start_rank_;

    PetscInt end_rank_;

    std::vector<PetscInt> sparsity_col_idx_;

    void create(const PETScSparsityPattern& sparsity_pattern);

    void preallocateFromSparsityPattern(
        const PETScSparsityPattern& sparsity_pattern);
 
    friend bool finalizeMatrixAssembly(PETScMatrix& mat,
                                       const MatAssemblyType asm_type);
};
 
template <class T_DENSE_MATRIX>
void PETScMatrix::add(std::vector<PetscInt> const& row_pos,
                      std::vector<PetscInt> const& col_pos,
                      const T_DENSE_MATRIX& sub_mat)
{
    const PetscInt nrows = static_cast<PetscInt>(row_pos.size());
    const PetscInt ncols = static_cast<PetscInt>(col_pos.size());
 
    PetscCallAbort(PETSC_COMM_WORLD,
                   MatSetValues(A_, nrows, &row_pos[0], ncols, &col_pos[0],
                                sub_mat.data(), ADD_VALUES));
};

bool finalizeMatrixAssembly(
    PETScMatrix& mat, const MatAssemblyType asm_type = MAT_FINAL_ASSEMBLY);

void setPreallocationNonzeroOption(PETScMatrix& matrix, bool has_col_idx);

MatStructure nonzeroPatternStructure(std::vector<PetscInt> const& col_a,
                                     std::vector<PetscInt> const& col_b);
 
}  // namespace MathLib
 
#include "MathLib/LinAlg/SetMatrixSparsity.h"
 
namespace MathLib
{
template <typename SPARSITY_PATTERN>
struct SetMatrixSparsity<PETScMatrix, SPARSITY_PATTERN>
{
    void operator()(PETScMatrix& matrix,
                    SPARSITY_PATTERN const& sparsity_pattern) const
    {
        setPreallocationNonzeroOption(matrix,
                                      !sparsity_pattern.col_idx.empty());
    }
};
}  // namespace MathLib