SparseMatrix_petsc.hpp

/*
 * SparseMatrix_petsc.hpp
 *
 *  Created on: Apr 26, 2016
 *      Author: i-bird
 */
 
#ifndef OPENFPM_NUMERICS_SRC_MATRIX_SPARSEMATRIX_PETSC_HPP_
#define OPENFPM_NUMERICS_SRC_MATRIX_SPARSEMATRIX_PETSC_HPP_
 
#include "util/petsc_util.hpp"
#include "Vector/map_vector.hpp"
#include <boost/mpl/int.hpp>
#include <petscmat.h>
#include "VTKWriter/VTKWriter.hpp"
#include "CSVWriter/CSVWriter.hpp"
 
#define PETSC_BASE 2
 
template<typename T>
class triplet<T,PETSC_BASE>
{
    PetscInt row_;
 
    PetscInt col_;
 
    PetscScalar val_;
 
public:
 
    PetscInt & row()
    {
        return row_;
    }
 
    PetscInt & col()
    {
        return col_;
    }
 
    PetscScalar & value()
    {
        return val_;
    }
 
    triplet(long int i, long int j, T val)
    {
        row_ = i;
        col_ = j;
        val_ = val;
    }
 
    // Default constructor
    triplet()
    :row_(0),col_(0),val_(0)
    {};
};
 
template<typename T, typename id_t>
class SparseMatrix<T,id_t, PETSC_BASE>
{
public:
 
    typedef boost::mpl::int_<PETSC_BASE> triplet_impl;
 
    typedef triplet<T,PETSC_BASE> triplet_type;
 
private:
 
    size_t g_row;
    size_t g_col;
 
    size_t l_row;
    size_t l_col;
 
    size_t start_row;
 
    bool m_created = false;
 
    Mat mat;
 
    openfpm::vector<triplet_type> trpl;
 
 
    mutable openfpm::vector<PetscScalar> vals;
    mutable openfpm::vector<PetscInt> cols;
    mutable openfpm::vector<PetscInt> d_nnz;
    mutable openfpm::vector<PetscInt> o_nnz;
 
    void fill_petsc()
    {
        d_nnz.resize(l_row);
        o_nnz.resize(l_row);
 
        d_nnz.fill(0);
        o_nnz.fill(0);
 
        // Here we explore every row to count how many non zero we have in the diagonal matrix part,
        // and the non diagonal matrix part, needed by MatMPIAIJSetPreallocation
 
        size_t i = 0;
 
        // Set the Matrix from triplet
        while (i < trpl.size())
        {
            PetscInt row = trpl.get(i).row();
 
            while(i < trpl.size() && row == trpl.get(i).row())
            {
                if ((size_t)trpl.get(i).col() >= start_row && (size_t)trpl.get(i).col() < start_row + l_row)
                    d_nnz.get(row - start_row)++;
                else
                    o_nnz.get(row - start_row)++;
                i++;
            }
        }
 
        PETSC_SAFE_CALL(MatMPIAIJSetPreallocation(mat,0,static_cast<const PetscInt*>(d_nnz.getPointer()),0,
                                                        static_cast<const PetscInt*>(o_nnz.getPointer())));
 
        // Counter i is zero
        i = 0;
 
        // Set the Matrix from triplet
        while (i < trpl.size())
        {
            vals.clear();
            cols.clear();
 
            PetscInt row = trpl.get(i).row();
 
            while(i < trpl.size() && row == trpl.get(i).row())
            {
                vals.add(trpl.get(i).value());
                cols.add(trpl.get(i).col());
                i++;
            }
            PETSC_SAFE_CALL(MatSetValues(mat,1,&row,cols.size(),static_cast<const PetscInt*>(cols.getPointer()),
                                                                static_cast<const PetscScalar *>(vals.getPointer()),
                                                                INSERT_VALUES));
        }
 
        PETSC_SAFE_CALL(MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY));
        PETSC_SAFE_CALL(MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY));
 
        m_created = true;
    }
 
    SparseMatrix(const SparseMatrix<T,id_t, PETSC_BASE> & spm)  {};
 
 
public:
 
    SparseMatrix(size_t N1, size_t N2, size_t n_row_local)
    :g_row(N1),g_col(N2),l_row(n_row_local),l_col(n_row_local)
    {
        PETSC_SAFE_CALL(MatCreate(PETSC_COMM_WORLD,&mat));
        PETSC_SAFE_CALL(MatSetType(mat,MATMPIAIJ));
        PETSC_SAFE_CALL(MatSetFromOptions(mat));
        PETSC_SAFE_CALL(MatSetSizes(mat,n_row_local,n_row_local,N1,N2));
 
        Vcluster<> & v_cl = create_vcluster();
 
        openfpm::vector<size_t> vn_row_local;
        v_cl.allGather(l_row,vn_row_local);
        v_cl.execute();
 
        // Calculate the starting row for this processor
 
        start_row = 0;
        for (size_t i = 0 ; i < v_cl.getProcessUnitID() ; i++)
            start_row += vn_row_local.get(i);
    }
 
    SparseMatrix()
    :g_row(0),g_col(0),l_row(0l),l_col(0),start_row(0)
    {
        PETSC_SAFE_CALL(MatCreate(PETSC_COMM_WORLD,&mat));
        PETSC_SAFE_CALL(MatSetType(mat,MATMPIAIJ));
        PETSC_SAFE_CALL(MatSetFromOptions(mat));
 
    }
 
    ~SparseMatrix()
    {
        // Destroy the matrix
        if (is_openfpm_init() == true)
        {PETSC_SAFE_CALL(MatDestroy(&mat));}
    }
 
    openfpm::vector<triplet_type> & getMatrixTriplets()
    {
        m_created = false;
 
        return this->trpl;
    }
 
    const Mat & getMat() const
    {
        if (m_created == false)
        {fill_petsc();}
 
        return mat;
    }
 
    Mat & getMat()
    {
        if (m_created == false)
        {fill_petsc();}
 
        return mat;
    }
 
    void resize(size_t row, size_t col, size_t l_row, size_t l_col)
    {
        if ((g_row != 0 && g_row != row) || (g_col != 0 && g_col != col) ||
            (this->l_row != 0 && this->l_row != l_row) || (this->l_col != 0 && this->l_col != l_col))
        {
            std::cout << __FILE__ << ":" << __LINE__ << " error you are resizing a PETSC matrix " << std::endl;
        }
 
        if (g_row != 0 && g_col != 0)   {return;}
 
        this->g_row = row;
        this->g_col = col;
 
        this->l_row = l_row;
        this->l_col = l_col;
 
        PETSC_SAFE_CALL(MatSetSizes(mat,l_row,l_col,g_row,g_col));
 
        Vcluster<> & v_cl = create_vcluster();
 
        openfpm::vector<size_t> vn_row_local;
        v_cl.allGather(l_row,vn_row_local);
        v_cl.execute();
 
        // Calculate the starting row for this processor
 
        start_row = 0;
        for (size_t i = 0 ; i < v_cl.getProcessUnitID() ; i++)
            start_row += vn_row_local.get(i);
    }
 
    T operator()(id_t i, id_t j)
    {
        T v;
 
        MatGetValues(mat,1,&i,1,&j,&v);
 
        return v;
    }
 
    T getValue(size_t r, size_t c)
    {
        for (size_t i = 0 ; i < trpl.size() ; i++)
        {
            if (r == (size_t)trpl.get(i).row() && c == (size_t)trpl.get(i).col())
                return trpl.get(i).value();
        }
 
        return 0;
    }
 
    bool isMatrixFilled()
    {
        return m_created;
    }
 
    /* Write matrix on vtk
     *
     * \param out file to write into
     *
     */
    bool write(std::string out,size_t opt = VTK_WRITER)
    {
        Vcluster<> & v_cl = create_vcluster();
 
        openfpm::vector<Point<2, double>> row_col;
        openfpm::vector<aggregate<double>> values;
 
        row_col.resize(trpl.size());
        values.resize(trpl.size());
 
        for (int i = 0 ; i < trpl.size() ; i++)
        {
            row_col.template get<0>(i)[1] = trpl.get(i).row();
            row_col.template get<0>(i)[0] = trpl.get(i).col();
 
            values.template get<0>(i) = trpl.get(i).value();
        }
 
        if (opt == VTK_WRITER)
        {
            auto ft = file_type::BINARY;
 
            // VTKWriter for a set of points
            VTKWriter<boost::mpl::pair<openfpm::vector<Point<2, double>>,
                                       openfpm::vector<aggregate<double>>>,
                                       VECTOR_POINTS> vtk_writer;
 
            vtk_writer.add(row_col,values,row_col.size());
 
            std::string output = std::to_string(out + "_" + std::to_string(v_cl.getProcessUnitID()) + std::to_string(".vtk"));
 
            openfpm::vector<std::string> prp_names;
            prp_names.add("value");
 
            // Write the VTK file
            return vtk_writer.write(output,prp_names,"matrix","",ft);
        }
        else
        {
            // CSVWriter test
            CSVWriter<openfpm::vector<Point<2,double>>,
                          openfpm::vector<aggregate<double>>> csv_writer;
 
            std::string output = std::to_string(out + "_" + std::to_string(v_cl.getProcessUnitID()) + std::to_string(".csv"));
 
            // Write the CSV
            return csv_writer.write(output,row_col,values);
        }
    }
};
 
 
#endif /* OPENFPM_NUMERICS_SRC_MATRIX_SPARSEMATRIX_PETSC_HPP_ */