array.hpp

#ifndef ARRAY_HPP_
#define ARRAY_HPP_

#include <vector>
#include <array>
#include <type_traits>
#include <initializer_list>
#include <stdexcept>


/*! Simple data-by-reference strided array class a la Fortran
 */
template <typename Scalar, size_t NDim, typename Storage = std::vector<Scalar> >
class Array
{
private:
    Storage& data_;
    std::array<size_t, NDim> shape_;
    std::array<size_t, NDim> stride_;
    size_t offset_ = 0;

public:
    Array(Storage& data, const std::array<size_t, NDim> shape, const std::array<size_t, NDim> stride, const size_t offset)
        : data_(data), shape_(shape), stride_(stride), offset_(offset)
        {}

    Array(Storage& data, const std::array<size_t, NDim> shape, const bool row_major=true)
        : data_(data), shape_(shape), offset_(0)
        {
            if (row_major) {
                for (size_t i = NDim; i > 0; --i) {
                    if (i == NDim)
                        stride_[i-1] = 1;
                    else
                        stride_[i-1] = shape[i] * stride_[i];
                }
            } else {
                for (size_t i = 0; i < NDim; ++i) {
                    if (i == 0)
                        stride_[i] = 1;
                    else
                        stride_[i] = shape[i-1] * stride_[i-1];
                }
            }

#ifndef NO_BOUNDS_CHECK
            size_t total = 1;
            for (size_t i = 0; i < NDim; ++i)
                total *= shape[i];
            if (total > data.size())
                throw std::out_of_range("data array too small");
#endif
        }

    template <typename... Idx>
    Scalar& operator()(Idx... idxs) { return data_[index(idxs...)]; }

    template <typename... Idx>
    const Scalar&& operator()(Idx... idxs) const { return data_[index(idxs...)]; }

    template <size_t axis>
    Array<Scalar, NDim-1> slice(size_t pos=0) const
        {
            static_assert(axis < NDim, "invalid axis");

            std::array<size_t, NDim-1> shape;
            std::array<size_t, NDim-1> stride;
            size_t offset;

            offset = offset_ + stride_[axis] * pos;

            for (size_t i = 0, j = 0; i < NDim; ++i) {
                if (i == axis)
                    continue;
                shape[j] = shape_[i];
                stride[j] = stride_[i];
                ++j;
            }

            return {data_, shape, stride, offset};
        }

    template <size_t axis>
    Array<Scalar, NDim> slice(size_t begin, size_t end) const
        {
            static_assert(axis < NDim, "invalid axis");

            std::array<size_t, NDim> shape;
            std::array<size_t, NDim> stride;
            size_t offset;

#ifndef NO_BOUNDS_CHECK
            if (end < begin || begin >= shape[axis] || end > shape[axis])
                throw std::out_of_range("begin/end indices out of bounds");
#endif

            offset = offset_ + stride_[axis] * begin;

            for (size_t i = 0; i < NDim; ++i) {
                if (i == axis)
                    shape[i] = end - begin;
                else
                    shape[i] = shape_[i];
                stride[i] = stride_[i];
            }

            return {data_, shape, stride, offset};
        }

    Storage& data() { return data_; }

    template <size_t axis>
    const size_t shape() const
        {
            static_assert(axis < NDim, "invalid axis");
            return shape_[axis];
        }

    template <size_t axis>
    const size_t stride() const
        {
            static_assert(axis < NDim, "invalid axis");
            return stride_[axis];
        }

    size_t offset() const { return offset_; }

    template <typename... Idx>
    size_t index(Idx... idxs) const
        {
            static_assert(sizeof...(idxs) == NDim,
                          "number of indices must equal the number of dimensions");

            const std::array<size_t, NDim> m{idxs...};
            size_t idx = offset_;

            for (size_t i = 0; i < NDim; ++i) {
                idx += stride_[i] * m[i];
#ifndef NO_BOUNDS_CHECK
                if (m[i] >= shape_[i])
                    throw std::out_of_range("index out of bounds");
#endif
            }

#ifndef NO_BOUNDS_CHECK
            if (idx >= data_.size())
                throw std::out_of_range("total index out of bounds");
#endif
            return idx;
        }

    typedef Eigen::Matrix<Scalar,Eigen::Dynamic,Eigen::Dynamic> EigenMatrix;
    typedef Eigen::Map<EigenMatrix, 0, Eigen::Stride<Eigen::Dynamic,Eigen::Dynamic> > EigenMap;

    EigenMap to_matrix() const
        {
            static_assert(NDim == 2, "matrix must be two-dimensional");
            return {(Scalar *)data_.data() + offset_,
                static_cast<Eigen::Index>(shape_[0]), static_cast<Eigen::Index>(shape_[1]),
                {static_cast<Eigen::Index>(stride_[1]), static_cast<Eigen::Index>(stride_[0])}};
        }

    operator EigenMap() const { return to_matrix(); }
};


#endif