MultiArray.h

// ===========================================================================
// Imagine++ Libraries
// Copyright (C) Imagine
// For detailed information: http://imagine.enpc.fr/software
// ===========================================================================

namespace Imagine {


    template <typename T, int dim> class MultiArray : public Array<T> {

    private:
        // Based on Array
        typedef Array<T> Base;
        Coords<dim> _sz;                    // Sizes of the array.
        FArray<size_t,dim> _stride;     // Strides to navigate to neighbors in each dimension
        // Set size and compute stride
        void setSizes(const Coords<dim>& sz)
        {
            _sz=sz;
            _stride[0] = 1;
            for (int i=0;i<dim-1;i++) {
                _stride[i+1] = _stride[i]*_sz[i];
            }
        }
        //  Set size and stride from another MultiArray
        template <typename Y> void setSizes(const MultiArray<Y,dim>& A)
        {
            _sz=A.sizes();  _stride=A.stride();
        }

    public:
        typedef T* iterator;
        typedef const T* const_iterator;


        MultiArray() : Base() { _sz = 0; }
        explicit MultiArray(const Coords<dim>& sz) : Base(sz.prod()) { setSizes(sz); }
        MultiArray(int s0,int s1) : Base(size_t(s0)*s1) { assert(dim==2); setSizes(Coords<2>(s0,s1)); }
        MultiArray(int s0,int s1, int s2) : Base(size_t(s0)*s1*s2) { assert(dim==3); setSizes(Coords<3>(s0,s1,s2)); }
        MultiArray(T* ptr, const Coords<dim>& sz,bool handleDelete=false) : Base(ptr,sz.prod(),handleDelete) { setSizes(sz); }
        MultiArray(T* ptr, int s0,int s1,bool handleDelete=false) : Base(ptr,size_t(s0)*s1,handleDelete) { assert(dim==2); setSizes(Coords<2>(s0,s1)); }
        MultiArray(T* ptr, int s0,int s1,int s2,bool handleDelete=false) : Base(ptr,size_t(s0)*s1*s2,handleDelete) { assert(dim==3); setSizes(Coords<3>(s0,s1,s2)); }
        explicit MultiArray(const MultiArray& A) : Base(A) { setSizes(A); }
        template <typename T2> 
        MultiArray(const MultiArray<T2,dim>& A) : Base(A) { setSizes(A); }
        virtual ~MultiArray() {}
        void setSize(const Coords<dim>& sz) {
            if (sz==_sz) return;
            Base::setSize(sz.prod());
            setSizes(sz);
        }
        void setSize(int s0, int s1) { setSize(Coords<2>(s0,s1)); }
        void setSize(int s0, int s1, int s2) { setSize(Coords<3>(s0,s1,s2)); }
        MultiArray& operator=(const MultiArray& A) { Base::operator =(A); setSizes(A); return *this; }
        template <typename T2> MultiArray& operator= (const MultiArray<T2,dim>& A) { Base::operator = (A); setSizes(A); return *this; }
        MultiArray clone() const {
            if (this->empty()) return MultiArray();
            MultiArray A(_sz);
            std::copy( this->begin(), this->end(), A.begin() );
            return A;
        }
        MultiArray& fill(T x) { Base::fill(x); return *this; }
        Coords<dim> sizes() const { return _sz; }
        size_t totalSize() const { return Array<T>::size(); }
        int size(int i) const { assert(i>=0 && i<dim); return _sz[i]; }
        int width() const { assert(dim>=1); return _sz[0]; }
        int height() const { assert(dim>=2); return _sz[1]; }
        int depth() const { assert(dim>=3); return _sz[2]; }
        FArray<size_t,dim> stride() const { return _stride; }
        size_t stride(int i) const { assert(i>=0 && i<dim); return _stride[i]; }
        MultiArray getSubArray(const Coords<dim>& offset, const Coords<dim>& sz) const {
            MultiArray A(sz);
            for (CoordsIterator<dim> p = A.coordsBegin() ; p != A.coordsEnd() ; ++p) {
                A(*p)=operator()(*p+offset);
            }
            return A;
        }
        size_t offset(const Coords<dim>& c) const {
            size_t o = c[dim-1];
            assert(o<size_t(_sz[dim-1]));
            for (int d=dim-2;d>=0;d--) {
                assert(c[d]>=0 && c[d]<_sz[d]);
                o=o*_sz[d]+c[d];
            }
            return o;
        }
        size_t offset(int x, int y) const {
            assert(dim==2);
            assert(x>=0); assert(x<_sz[0]);
            assert(y>=0); assert(y<_sz[1]);
            return x+(size_t)_sz[0]*y;
        }
        size_t offset(int x, int y, int z) const {
            assert(dim==3);
            assert(x>=0); assert(x<_sz[0]);
            assert(y>=0); assert(y<_sz[1]);
            assert(z>=0); assert(z<_sz[2]);
            return x+_sz[0]*(y+(size_t)_sz[1]*z);
        }
        const T& operator() (const Coords<dim>& c) const {
            return (*this)[offset(c)];
        }
        T& operator() (const Coords<dim>& c) { return (*this)[offset(c)]; }
        const T& operator()(int x, int y) const { return (*this)[offset(x,y)]; }
        T& operator()(int x, int y) { return (*this)[offset(x,y)]; }
        const T& operator()(int x, int y, int z) const { return (*this)[offset(x,y,z)]; } 
        T& operator()(int x, int y, int z) { return (*this)[offset(x,y,z)]; }
        CoordsIterator<dim> coordsBegin() const { return CoordsIterator<dim>( Coords<dim>(0), _sz -Coords<dim>(1) ); }
        CoordsIterator<dim> coordsEnd() const { return CoordsIterator<dim>(); }
        friend void write(std::ostream& out,const MultiArray& A) {
            write(out,A._sz);
            out.write((const char*)A.data(),A.totalSize()*sizeof(T));
        }

        friend void read(std::istream& in,MultiArray& A) {
            Coords<dim> sz;
            read(in,sz);
            A.setSize(sz);
            in.read((char*)A.data(),A.totalSize()*sizeof(T));
        }
        friend inline std::ostream& operator<<(std::ostream& out,const MultiArray& A) {
            out<<A.sizes()<<" ";
            for (size_t i=0;i<A.totalSize();i++) {
                out<<A[i];
                if (i<A.totalSize()-1) out<<" ";
            }
            return out;
        }
        friend inline std::istream& operator>>(std::istream& in,MultiArray& A) {
            Coords<dim> sz;
            in>>sz;
            A.setSize(sz);
            for (size_t i=0;i<A.totalSize();i++) 
                in>>A[i];
            return in;
        }
    };


}