PETScVector.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 <petscvec.h>
 
#include <map>
#include <span>
#include <string>
#include <vector>
 
namespace MathLib
{
class PETScVector
{
public:
    using IndexType = PetscInt;
    // TODO make this class opaque, s.t. the PETSc symbols are not scattered all
    //      over the global namespace
    using PETSc_Vec = Vec;
 
public:
    PETScVector() {}
    PETScVector(const PetscInt vec_size, const bool is_global_size = true);

    PETScVector(const PetscInt vec_size, const std::vector<PetscInt>& ghost_ids,
                const bool is_global_size = true);

    explicit PETScVector(const PETScVector& existing_vec,
                         const bool deep_copy = true);
 
    PETScVector(PETScVector&& other);
 
    ~PETScVector() { destroy(); }
    void finalizeAssembly();

    PetscInt size() const { return size_; }
    PetscInt getLocalSize() const { return size_loc_; }
    PetscInt getGhostSize() const { return size_ghosts_; }
    PetscInt getRangeBegin() const { return start_rank_; }
    PetscInt getRangeEnd() const { return end_rank_; }
    void set(const PetscInt i, const PetscScalar value)
    {
        PetscCallAbort(PETSC_COMM_WORLD,
                       VecSetValue(v_, i, value, INSERT_VALUES));
    }

    void add(const PetscInt i, const PetscScalar value)
    {
        PetscCallAbort(PETSC_COMM_WORLD, VecSetValue(v_, i, value, ADD_VALUES));
    }

    template <class T_SUBVEC>
    void add(const std::vector<PetscInt>& e_idxs, const T_SUBVEC& sub_vec)
    {
        if constexpr (std::is_pointer_v<T_SUBVEC>)
        {
            PetscCallAbort(PETSC_COMM_WORLD,
                           VecSetValues(v_, e_idxs.size(), &e_idxs[0],
                                        &sub_vec[0], ADD_VALUES));
        }
        else
        {
            PetscCallAbort(PETSC_COMM_WORLD,
                           VecSetValues(v_, e_idxs.size(), e_idxs.data(),
                                        sub_vec.data(), ADD_VALUES));
        }
    }

    template <class T_SUBVEC>
    void set(const std::vector<PetscInt>& e_idxs, const T_SUBVEC& sub_vec)
    {
        if constexpr (std::is_pointer_v<T_SUBVEC>)
        {
            PetscCallAbort(PETSC_COMM_WORLD,
                           VecSetValues(v_, e_idxs.size(), e_idxs.data(),
                                        sub_vec, INSERT_VALUES));
        }
        else
        {
            PetscCallAbort(PETSC_COMM_WORLD,
                           VecSetValues(v_, e_idxs.size(), e_idxs.data(),
                                        sub_vec.data(), INSERT_VALUES));
        }
    }
 
    // TODO preliminary
    void setZero() { PetscCallAbort(PETSC_COMM_WORLD, VecSet(v_, 0.0)); }
    PETScVector& operator=(PETScVector&&) = delete;

    void setLocalAccessibleVector() const;

    std::vector<PetscScalar> get(std::vector<IndexType> const& indices) const;

    PetscScalar operator[](PetscInt idx) const { return get(idx); }
    void getGlobalVector(std::vector<PetscScalar>& u) const;
 
    /* Get an entry value. This is an expensive operation,
       and it only get local value. Use it for only test purpose
       Get the value of an entry by [] operator.
       setLocalAccessibleVector() must be called beforehand.
    */
    PetscScalar get(const PetscInt idx) const;

    PETSc_Vec& getRawVector() { return v_; }
    PETSc_Vec const& getRawVector() const { return v_; }
    void copyValues(std::vector<PetscScalar>& u) const;

    void copyValues(std::span<PetscScalar> u) const;

    void viewer(
        const std::string& file_name,
        const PetscViewerFormat vw_format = PETSC_VIEWER_ASCII_MATLAB) const;
 
    void shallowCopy(const PETScVector& v);
 
private:
    void destroy()
    {
        if (v_ != nullptr)
        {
            PetscCallAbort(PETSC_COMM_WORLD, VecDestroy(&v_));
        }
        v_ = nullptr;
    }
 
    PETSc_Vec v_ = nullptr;
    mutable PETSc_Vec v_loc_ = nullptr;

    PetscInt start_rank_;
    PetscInt end_rank_;

    PetscInt size_;
    PetscInt size_loc_;
    PetscInt size_ghosts_ = 0;

    bool created_with_ghost_id_ = false;

    mutable std::vector<PetscScalar> entry_array_;

    mutable std::map<PetscInt, PetscInt> global_ids2local_ids_ghost_;

    PetscScalar* getLocalVector() const;

    PetscInt getLocalIndex(const PetscInt global_index) const;

    inline void restoreArray(PetscScalar* array) const;

    void config();
 
    friend void finalizeVectorAssembly(PETScVector& vec);
};

void finalizeVectorAssembly(PETScVector& vec);
 
}  // namespace MathLib