PETScVector.h

Go to the documentation of this file.

 
#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)
    {
        VecSetValue(v_, i, value, INSERT_VALUES);
    }
 
    void add(const PetscInt i, const PetscScalar value)
    {
        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>)
        {
            VecSetValues(v_, e_idxs.size(), e_idxs.data(), sub_vec, ADD_VALUES);
        }
        else
        {
            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>)
        {
            VecSetValues(v_, e_idxs.size(), e_idxs.data(), sub_vec,
                         INSERT_VALUES);
        }
        else
        {
            VecSetValues(v_, e_idxs.size(), e_idxs.data(), sub_vec.data(),
                         INSERT_VALUES);
        }
    }
 
    // TODO preliminary
    void setZero() { 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)
        {
            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