ndarray.h

Go to the documentation of this file.

 
#ifndef NDARRAY_H
#define NDARRAY_H
 
#include <vector>
#include <cassert>
#include <iostream>
#include <algorithm>
#include <type_traits>
#include <numeric>
// for heterogeneous computing, we can use ATen::Tensor
//#include "./module_container/ATen/tensor.h"
 
template<typename T>
class NDArray
{
    // align with STL container implementation, there are several functions compulsory to be implemented
    // constructor: default, copy, move, initializer_list
    // operator: =, ==, !=, <, <=, >, >=
    // iterator: begin, cbegin, end, cend
    // capacity: size, empty, max_size, reserve, shrink_to_fit
    // element access: [], at, front, back, data
    // modifiers: clear, insert, erase, push_back, pop_back, resize, swap
    // allocator: get_allocator
public:
    // constructors
    NDArray()= delete;
    // initializer_list constructor
    NDArray(std::initializer_list<size_t> il) : shape_(il), data_(std::accumulate(shape_.begin(), shape_.end(), 1, std::multiplies<size_t>())) {}
    NDArray(std::initializer_list<int> il) : shape_(il.begin(), il.end()), data_(std::accumulate(shape_.begin(), shape_.end(), 1, std::multiplies<size_t>())) {}
    // variadic template constructor, (delegate constructor)
    template<typename... Args> NDArray(const size_t idx, Args... args) : NDArray({idx, static_cast<size_t>(args)...}) {}
    template<typename... Args> NDArray(const int& idx, Args... args) : NDArray({idx, static_cast<int>(args)...}) {} // not happy with this because size_t can have larger range
    // copy constructor
    NDArray(const NDArray& other) : data_(other.data_), shape_(other.shape_) {}
    // move constructor
    NDArray(NDArray&& other) : data_(std::move(other.data_)), shape_(std::move(other.shape_)) {}
 
    // destructor
    ~NDArray() {}
 
    // operators
    NDArray& operator=(const NDArray& other)
    {
        if (this != &other)
        {
            data_ = other.data_;
            shape_ = other.shape_;
        }
        return *this;
    }
    NDArray& operator=(NDArray&& other)
    {
        if (this != &other)
        {
            data_ = std::move(other.data_);
            shape_ = std::move(other.shape_);
        }
        return *this;
    }
 
    bool operator==(const NDArray& other) const { return data_ == other.data_ && shape_ == other.shape_; }
    bool operator!=(const NDArray& other) const { return !(*this == other); }
    // other operators are not generally supported
 
    // element access
    template<typename... Args> T& at(const size_t idx, Args... args) { return data_[index(idx, args...)]; }
    template<typename... Args> const T& at(const size_t idx, Args... args) const { return data_[index(idx, args...)]; }
    template<typename... Args> T& operator()(const size_t idx, Args... args) { return data_[index(idx, args...)]; }
    template<typename... Args> const T& operator()(const size_t idx, Args... args) const { return data_[index(idx, args...)]; }
    // front
    T& front() { return data_.front(); }
    const T& front() const { return data_.front(); }
    // back
    T& back() { return data_.back(); }
    const T& back() const { return data_.back(); }
    // data
    T* data() { return data_.data(); }
    const T* data() const { return data_.data(); }
 
    // iterators
    // iterators on the whole data
    T* begin() { return data_.data(); }
    T* end() { return data_.data() + data_.size(); }
    const T* cbegin() const { return data_.data(); }
    const T* cend() const { return data_.data() + data_.size(); }
    // iterators on different dimensions
    
    // capacity
    // size
    size_t size() const { return data_.size(); }
    size_t size(const size_t& dim) const { return shape_.at(dim); }
    // empty
    bool empty() const { return data_.empty(); }
    // multi-dimensional specific
    // shape
    const std::vector<size_t>& shape() const { return shape_; }
    // reshape
    template<typename... Args>
    void reshape(Args... args)
    {
        // DEVELP WARNING: what if arg = -2? :)
        // save args into a vector
        //std::vector<int64_t> dims = {static_cast<int64_t>(args)...};
        std::vector<int64_t> dims = {args...};
        // assert number of -1 in dims is at most 1
        // -1 is not type-safe!!!
        size_t count = std::count_if(dims.begin(), dims.end(), [](size_t i) { return i == -1; });
        assert(count <= 1);
        // if there is -1, calculate the size
        if (count == 1)
        {
            size_t size = 1;
            for (size_t i = 0; i < dims.size(); ++i)
            {
                if (dims[i] != -1)
                {
                    size *= dims[i];
                }
            }
            size_t idx = std::find(dims.begin(), dims.end(), -1) - dims.begin();
            dims[idx] = data_.size() / size;
        }
        // calculate the size
        size_t size = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<size_t>());
        // assert size is the same
        assert(size == data_.size());
        // assign dims to shape_
        std::copy(dims.begin(), dims.end(), shape_.begin());
    }
 
    // interface to ATen::Tensor, but constraint to int, double, float, std::complex<float>, std::complex<double>
    // std::enable_if<
    //     std::is_same<T, int>::value 
    //  || std::is_same<T, double>::value 
    //  || std::is_same<T, float>::value 
    //  || std::is_same<T, std::complex<float>>::value 
    //  || std::is_same<T, std::complex<double>>::value, container::Tensor
    // >::type to_tensor() const
    // {
    //     container::TensorShape shape(shape_);
    //     container::Tensor result = container::Tensor(container::DataTypeToEnum<T>::value, shape);
    //     std::memcpy(result.data<T>(), data_.data(), data_.size() * sizeof(T));
    //     return result;
    // }
    template<typename... Args>
    size_t index(const size_t idx, Args... args) const
    {
        assert(sizeof...(args) == shape_.size() - 1); // assert the indices are the same as the shape
        size_t indices[] = {idx, static_cast<size_t>(args)...};
        size_t index = 0;
        for (size_t i = 0; i < shape_.size(); ++i)
        {
            index += indices[i] * std::accumulate(shape_.begin() + i + 1, shape_.end(), 1, std::multiplies<size_t>());
        }
        assert(index < data_.size()); // assert the index is within the data
        return index;
    }
private:
    std::vector<size_t> shape_;
    // for GPU-compatible data container, will be replaced by raw pointer
    std::vector<T> data_;
};
 
#endif // NDARRAY_H