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>

#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;
    }


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(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));
    }

    ~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)
    {
        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;
    }
};


#endif /* OPENFPM_NUMERICS_SRC_MATRIX_SPARSEMATRIX_PETSC_HPP_ */