Vector_petsc.hpp

/*
 * Vector_petsc.hpp
 *
 *  Created on: Apr 29, 2016
 *      Author: i-bird
 */

#ifndef OPENFPM_NUMERICS_SRC_VECTOR_VECTOR_PETSC_HPP_
#define OPENFPM_NUMERICS_SRC_VECTOR_VECTOR_PETSC_HPP_

#include "Vector/map_vector.hpp"
#include "Vector/vector_def.hpp"
#include <boost/mpl/int.hpp>
#include <petscvec.h>
#include "util/petsc_util.hpp"
#include <unordered_map>

#define PETSC_RVAL 2

template<typename T>
class rval<T,PETSC_RVAL>
{
public:

    typedef boost::fusion::vector<PetscInt,T> type;

    type data;

    static const unsigned int row = 0;

    static const unsigned int value = 1;

    static const unsigned int max_prop = 2;

    long int & rw()
    {
        return boost::fusion::at_c<row>(data);
    }

    T & val()
    {
        return boost::fusion::at_c<value>(data);
    }

    rval()  {}

    rval(long int i, T val)
    {
        rw() = i;
        val() = val;
    }

    static inline bool noPointers()
    {
        return true;
    }
};

constexpr unsigned int row_id = 0;
constexpr unsigned int val_id = 1;


template<typename T>
class Vector<T,PETSC_BASE>
{
    size_t n_row;

    size_t n_row_local;

    mutable bool v_created = false;

    mutable Vec v;

    mutable openfpm::vector<rval<PetscScalar,PETSC_RVAL>,HeapMemory,typename memory_traits_inte<rval<PetscScalar,PETSC_RVAL>>::type, memory_traits_inte > row_val;

    mutable std::unordered_map<size_t,size_t> map;

    T invalid;

    void setPetsc() const
    {
        if (v_created == false)
        {PETSC_SAFE_CALL(VecSetType(v,VECMPI));}


        // set the vector

        if (row_val.size() != 0)
            PETSC_SAFE_CALL(VecSetValues(v,row_val.size(),&row_val.template get<row_id>(0),&row_val.template get<val_id>(0),INSERT_VALUES))

        PETSC_SAFE_CALL(VecAssemblyBegin(v));
        PETSC_SAFE_CALL(VecAssemblyEnd(v));

        v_created = true;
    }

public:

    Vector(const Vector<T,PETSC_BASE> & v)
    {
        this->operator=(v);
    }

    Vector(Vector<T,PETSC_BASE> && v)
    :n_row(0),n_row_local(0),invalid(0)
    {
        this->operator=(v);
    }

    ~Vector()
    {
        if (is_openfpm_init() == true)
        {PETSC_SAFE_CALL(VecDestroy(&v));}
    }

    Vector(size_t n, size_t n_row_local)
    :n_row_local(n_row_local),v(NULL),invalid(0)
    {
        // Create the vector
        PETSC_SAFE_CALL(VecCreate(PETSC_COMM_WORLD,&v));

        resize(n,n_row_local);
    }

    Vector()
    :n_row(0),n_row_local(0),invalid(0)
    {
        // Create the vector
        PETSC_SAFE_CALL(VecCreate(PETSC_COMM_WORLD,&v));
    }

    void resize(size_t row, size_t l_row)
    {
        n_row = row;
        n_row_local = l_row;

        PETSC_SAFE_CALL(VecSetSizes(v,n_row_local,n_row));
    }

    void insert(size_t i, T val)
    {
        row_val.add();

        // Map
        map[i] = row_val.size()-1;

        row_val.last().template get<row_id>() = i;
        row_val.last().template get<val_id>() = val;
    }

    inline PetscScalar & insert(size_t i)
    {
        row_val.add();

        // Map
        map[i] = row_val.size()-1;

        row_val.last().template get<row_id>() = i;
        return row_val.last().template get<val_id>();
    }

    inline const PetscScalar & insert(size_t i) const
    {
        row_val.add();

        // Map
        map[i] = row_val.size()-1;

        row_val.last().template get<row_id>() = i;
        return row_val.last().template get<val_id>();
    }

    const PetscScalar & operator()(size_t i) const
    {
        // Search if exist

        std::unordered_map<size_t,size_t>::iterator it = map.find(i);

        if ( it != map.end() )
            return row_val.template get<val_id>(it->second);

        return insert(i);
    }

    PetscScalar & operator()(size_t i)
    {
        // Search if exist

        std::unordered_map<size_t,size_t>::iterator it = map.find(i);

        if ( it != map.end() )
            return row_val.template get<val_id>(it->second);

        return insert(i);
    }

    const Vec & getVec() const
    {
        setPetsc();

        return v;
    }

    Vec & getVec()
    {
        setPetsc();

        return v;
    }

    void update()
    {
        PetscInt n_row;
        PetscInt n_row_local;

        // Get the size of the vector from PETSC
        VecGetSize(v,&n_row);
        VecGetLocalSize(v,&n_row_local);

        this->n_row = n_row;
        this->n_row_local = n_row_local;

        row_val.resize(n_row_local);

        //

        PetscInt low;
        PetscInt high;

        VecGetOwnershipRange(v,&low,&high);

        // Fill the index and construct the map

        size_t k = 0;
        for (size_t i = low ; i < (size_t)high ; i++)
        {
            row_val.template get<row_id>(k) = i;
            map[i] = k;
            k++;
        }

        PETSC_SAFE_CALL(VecGetValues(v,row_val.size(),&row_val.template get<row_id>(0),&row_val.template get<val_id>(0)))
    }

    Vector<T,PETSC_BASE> & operator=(const Vector<T,PETSC_BASE> & v)
    {
        map = v.map;
        row_val = v.row_val;

        return *this;
    }

    Vector<T,PETSC_BASE> & operator=(Vector<T,PETSC_BASE> && v)
    {
        map.swap(v.map);
        row_val.swap(v.row_val);

        return *this;
    }

    void setZero()
    {
        if (v_created == false)
        {PETSC_SAFE_CALL(VecSetType(v,VECMPI));}

        v_created = true;
    }
};


#endif /* OPENFPM_NUMERICS_SRC_VECTOR_VECTOR_EIGEN_HPP_ */