spm.h

#ifndef gridripper_math_spm_h
#define gridripper_math_spm_h

#include <cstring>

namespace gridripper { namespace math {

template <class TD>
class SpM
{
private:
    int numRows;
    int numColumns;

protected:
    SpM(int nrows, int ncols): numRows(nrows), numColumns(ncols) {
    }

    void resize(int nrows, int ncols) {
        numRows = nrows;
        numColumns = ncols;
    }

public:
    int getNumRows() const {
        return numRows;
    }

    int getNumColumns() const {
        return numColumns;
    }

    virtual int countNonzerosInRow(int row) const =0;

    virtual void gmv(TD* y, TD alpha, const TD* x, TD beta, const TD* b)
        const =0;

    virtual void toDenseMatrix(TD* mat) const =0;
};

template <class TD, class TI>
class DiaSpM: public SpM<TD>
{
private:
    TD* data;
    TI* offsets;
    int numDiagonals;

public:
    DiaSpM(int nrows, int ncols, int ndia): SpM<TD>(nrows, ncols), numDiagonals(ndia) {
        int n = nrows*ndia;
        data = new TD[n];
        offsets = new TI[ndia];
        for(int i = 0; i < n; ++i) {
            data[i] = 0;
        }
        for(int i = 0; i < ndia; ++i) {
            offsets[i] = 0;
        }
    }

    ~DiaSpM() {
        delete [] data;
        delete [] offsets;
    }

    void realloc(int nrows, int ncols, int ndia) {
        if(nrows == SpM<TD>::getNumRows()
                && ncols == SpM<TD>::getNumColumns()
                && numDiagonals == ndia) {
            return;
        }
        delete [] data;
        delete [] offsets;
        SpM<TD>::resize(nrows, ncols);
        numDiagonals = ndia;
        int n = nrows*ndia;
        data = new TD[n];
        offsets = new TI[ndia];
    }

    void setOffset(int dia, TI offset) {
        offsets[dia] = offset;
    }

    void set(int row, int dia, TD v) {
        data[row*numDiagonals + dia] = v;
    }

    const TD* internalArray() const {
        return data;
    }

    TD* internalArray() {
        return data;
    }

    int arrayLength() const {
        return SpM<TD>::getNumRows()*numDiagonals;
    }

    const TI* internalOffsetArray() const {
        return offsets;
    }

    TI* internalOffsetArray() {
        return offsets;
    }

    int getNumDiagonals() const {
        return numDiagonals;
    }

    int countNonzerosInRow(int row) const {
        int nrows = SpM<TD>::getNumRows();
        int ncols = SpM<TD>::getNumColumns();
        int n = 0;
        for(int dia = 0; dia < numDiagonals; ++dia) {
            int offset = offsets[dia];
            int start = offset < 0? -offset : 0;
            int end = ncols - offset;
            if(end > nrows) {
                end = nrows;
            }
            if(row >= start && row < end) {
                ++n;
            }
        }
        return n;
    }

    void gmv(TD* y, TD alpha, const TD* x, TD beta, const TD* b) const {
        int nrows = SpM<TD>::getNumRows();
        int ncols = SpM<TD>::getNumColumns();
        for(int row = 0; row < nrows; ++row) {
            TD v = 0;
            for(int dia = 0; dia < numDiagonals; ++dia) {
                int offset = offsets[dia];
                int start = offset < 0? -offset : 0;
                int end = ncols - offset;
                if(end > nrows) {
                    end = nrows;
                }
                if(row >= start && row < end) {
                    v += data[row*numDiagonals + dia]*x[row + offset];
                }
            }
            if(b != NULL) {
                y[row] = alpha*v + beta*b[row];
            } else {
                y[row] = alpha*v;
            }
        }
    }

    void toDenseMatrix(TD* mat) const {
        int nrows = SpM<TD>::getNumRows();
        int ncols = SpM<TD>::getNumColumns();
        memset(mat, 0, nrows*ncols*sizeof(mat[0]));
        for(int dia = 0; dia < numDiagonals; ++dia) {
            int offset = offsets[dia];
            int row = offset < 0? -offset : 0;
            int end = ncols - offset;
            if(end > nrows) {
                end = nrows;
            }
            while(row < end) {
                mat[row*ncols + row + offset] = data[row*numDiagonals + dia];
                ++row;
            }
        }
    }
};

template <class TD, class TI>
class EllSpM: public SpM<TD>
{
private:
    TD* data;
    TI* columnIndices;
    int maxNonzerosPerRow;

public:
    EllSpM(int nrows, int ncols, int nonzero): SpM<TD>(nrows, ncols),
            maxNonzerosPerRow(nonzero) {
        int n = nrows*nonzero;
        data = new TD[n];
        columnIndices = new TI[n];
        for(int i = 0; i < n; ++i) {
            data[i] = 0;
            columnIndices[i] = -1;
        }
    }

    ~EllSpM() {
        delete [] data;
        delete [] columnIndices;
    }

    void realloc(int nrows, int ncols, int nonzero) {
        if(nrows == SpM<TD>::getNumRows()
                && ncols == SpM<TD>::getNumColumns()
                && maxNonzerosPerRow == nonzero) {
            return;
        }
        delete [] data;
        delete [] columnIndices;
        SpM<TD>::resize(nrows, ncols);
        maxNonzerosPerRow = nonzero;
        int n = nrows*nonzero;
        data = new TD[n];
        columnIndices = new TI[n];
        for(int i = 0; i < n; ++i) {
            columnIndices[i] = -1;
        }
    }

    void set(int row, int col, TD v, int i) {
        int k = row*maxNonzerosPerRow + i;
        columnIndices[k] = col;
        data[k] = v;
    }

    const TD* internalArray() const {
        return data;
    }

    TD* internalArray() {
        return data;
    }

    const TI* internalColumnIndexArray() const {
        return columnIndices;
    }

    TI* internalColumnIndexArray() {
        return columnIndices;
    }

    int arrayLength() const {
        return SpM<TD>::getNumRows()*maxNonzerosPerRow;
    }

    int countNonzerosInRow(int row) const {
        int n = 0;
        int offset = row*maxNonzerosPerRow;
        for(int i = 0; i < maxNonzerosPerRow; ++i) {
            int icol = columnIndices[offset + i];
            if(icol < 0) {
                break;
            }
            ++n;
        }
        return n;
    }

    void gmv(TD* y, TD alpha, const TD* x, TD beta, const TD* b) const {
        int nrows = SpM<TD>::getNumRows();
        for(int irow = 0; irow < nrows; ++irow) {
            TD v = 0;
            int offset = irow*maxNonzerosPerRow;
            for(int i = 0; i < maxNonzerosPerRow; ++i) {
                int icol = columnIndices[offset + i];
                if(icol < 0) {
                    break;
                }
                v += data[offset + i]*x[icol];
            }
            if(b != NULL) {
                y[irow] = alpha*v + beta*b[irow];
            } else {
                y[irow] = alpha*v;
            }
        }
    }

    int getMaxNonzerosPerRow() const {
        return maxNonzerosPerRow;
    }

    void toDenseMatrix(TD* mat) const {
        int nrows = SpM<TD>::getNumRows();
        int ncols = SpM<TD>::getNumColumns();
        memset(mat, 0, nrows*ncols*sizeof(mat[0]));
        int offset = 0;
        for(int irow = 0; irow < nrows; ++irow) {
            for(int i = 0; i < maxNonzerosPerRow; ++i) {
                int icol = columnIndices[offset + i];
                if(icol < 0) {
                    break;
                }
                mat[irow*ncols + icol] = data[offset + i];
            }
            offset += maxNonzerosPerRow;
        }
    }
};

template <class TD, class TI>
class CooSpM: public SpM<TD>
{
private:
    TD* data;
    TI* rowIndices;
    TI* columnIndices;
    int nonzeroCount;

public:
    CooSpM(int nrows, int ncols, int maxnonzero): SpM<TD>(nrows, ncols),
            nonzeroCount(maxnonzero) {
        data = new TD[maxnonzero];
        rowIndices = new TI[maxnonzero];
        columnIndices = new TI[maxnonzero];
        for(int i = 0; i < maxnonzero; ++i) {
            data[i] = 0;
            rowIndices[i] = 0;
            columnIndices[i] = 0;
        }
    }

    ~CooSpM() {
        delete [] data;
        delete [] rowIndices;
        delete [] columnIndices;
    }

    void realloc(int nrows, int ncols, int nonzero) {
        if(nrows == SpM<TD>::getNumRows()
                && ncols == SpM<TD>::getNumColumns()
                && nonzero == nonzeroCount) {
            return;
        }
        delete [] data;
        delete [] rowIndices;
        delete [] columnIndices;
        SpM<TD>::resize(nrows, ncols);
        nonzeroCount = nonzero;
        data = new TD[nonzero];
        rowIndices = new TI[nonzero];
        columnIndices = new TI[nonzero];
    }

    void set(int row, int col, TD v, int i) {
        data[i] = v;
        rowIndices[i] = row;
        columnIndices[i] = col;
    }

    const TD* internalArray() const {
        return data;
    }

    TD* internalArray() {
        return data;
    }

    int arrayLength() const {
        return nonzeroCount;
    }

    const TI* internalRowIndexArray() const {
        return rowIndices;
    }

    TI* internalRowIndexArray() {
        return rowIndices;
    }

    const TI* internalColumnIndexArray() const {
        return columnIndices;
    }

    TI* internalColumnIndexArray() {
        return columnIndices;
    }

    int countNonzerosInRow(int row) const {
        int n = 0;
        for(int i = 0; i < nonzeroCount; ++i) {
            if(rowIndices[i] == row) {
                ++n;
            }
        }
        return n;
    }

    void gmv(TD* y, TD alpha, const TD* x, TD beta, const TD* b) const {
        int nrows = SpM<TD>::getNumRows();
        if(b != NULL) {
            for(int row = 0; row < nrows; ++row) {
                y[row] = beta*b[row];
            }
        } else {
            for(int row = 0; row < nrows; ++row) {
                y[row] = 0;
            }
        }
        for(int i = 0; i < nonzeroCount; ++i) {
            int irow = rowIndices[i];
            int icol = columnIndices[i];
            y[irow] += alpha*data[i]*x[icol];
        }
    }

    int getNonzeroCount() const {
        return nonzeroCount;
    }

    void toDenseMatrix(TD* mat) const {
        int nrows = SpM<TD>::getNumRows();
        int ncols = SpM<TD>::getNumColumns();
        memset(mat, 0, nrows*ncols*sizeof(mat[0]));
        for(int i = 0; i < nonzeroCount; ++i) {
            int irow = rowIndices[i];
            int icol = columnIndices[i];
            mat[irow*ncols + icol] = data[i];
        }
    }
};

template <class TD, class TI>
class CsrSpM: public SpM<TD>
{
private:
    TD* data;
    TI* columnIndices;
    TI* rowPtr;
    int nonzeroCount;

public:
    CsrSpM(int nrows, int ncols, int nonzero): SpM<TD>(nrows, ncols),
            nonzeroCount(nonzero) {
        data = new TD[nonzero];
        columnIndices = new TI[nonzero];
        rowPtr = new TI[nrows + 1];
        for(int i = 0; i < nonzero; ++i) {
            data[i] = 0;
            columnIndices[i] = 0;
        }
        for(int i = 0; i <= nrows; ++i) {
            rowPtr[i] = 0;
        }
    }

    ~CsrSpM() {
        delete [] data;
        delete [] columnIndices;
        delete [] rowPtr;
    }

    void realloc(int nrows, int ncols, int nonzero) {
        if(nrows == SpM<TD>::getNumRows()
                && ncols == SpM<TD>::getNumColumns()
                && nonzero == nonzeroCount) {
            return;
        }
        delete [] data;
        delete [] columnIndices;
        delete [] rowPtr;
        SpM<TD>::resize(nrows, ncols);
        nonzeroCount = nonzero;
        data = new TD[nonzero];
        columnIndices = new TI[nonzero];
        rowPtr = new TI[nrows + 1];
    }

    void setRowPtr(int row, TI p) {
        rowPtr[row] = p;
    }

    void set(int row, int col, TD v, int i) {
        data[i] = v;
        columnIndices[i] = col;
    }

    const TD* internalArray() const {
        return data;
    }

    TD* internalArray() {
        return data;
    }

    int arrayLength() const {
        return nonzeroCount;
    }

    const TI* internalRowPtrArray() const {
        return rowPtr;
    }

    TI* internalRowPtrArray() {
        return rowPtr;
    }

    const TI* internalColumnIndexArray() const {
        return columnIndices;
    }

    TI* internalColumnIndexArray() {
        return columnIndices;
    }

    int countNonzerosInRow(int row) const {
        return rowPtr[row + 1] - rowPtr[row];
    }

    void gmv(TD* y, TD alpha, const TD* x, TD beta, const TD* b) const {
        int nrows = SpM<TD>::getNumRows();
        if(b != NULL) {
            for(int row = 0; row < nrows; ++row) {
                y[row] = beta*b[row];
            }
        } else {
            for(int row = 0; row < nrows; ++row) {
                y[row] = 0;
            }
        }
        for(int irow = 0; irow < nrows; ++irow) {
            int start = rowPtr[irow];
            int end = rowPtr[irow + 1];
            for(int i = start; i < end; ++i) {
                y[irow] += alpha*data[i]*x[columnIndices[i]];
            }
        }
    }

    int getNonzeroCount() const {
        return nonzeroCount;
    }

    void toDenseMatrix(TD* mat) const {
        int nrows = SpM<TD>::getNumRows();
        int ncols = SpM<TD>::getNumColumns();
        memset(mat, 0, nrows*ncols*sizeof(mat[0]));
        for(int irow = 0; irow < nrows; ++irow) {
            int p = rowPtr[irow];
            int end = rowPtr[irow + 1];
            while(p < end) {
                mat[irow*ncols + columnIndices[p]] = data[p];
                ++p;
            }
        }
    }
};

} } // namespace gridripper::math

#endif /* gridripper_math_spm_h */