array_iterator.hpp

Go to the documentation of this file.

// mds_utils/python/numpy/array_iterator.hpp
//
// Copyright (c) 2014 - Michele De Stefano (micdestefano@users.sourceforge.net)
//
// Distributed under the MIT License (See accompanying file LICENSE)

#ifndef MDS_UTILS_PYTHON_NUMPY_ARRAY_ITERATOR_HPP_INCLUDED
#define MDS_UTILS_PYTHON_NUMPY_ARRAY_ITERATOR_HPP_INCLUDED

#include <boost/iterator/iterator_facade.hpp>
#include <numpy/arrayobject.h>
#include <mds_utils/python/numpy/traits.hpp>
#include <complex>
#include <type_traits>

namespace mds_utils {
    namespace python {

        namespace numpy {

struct c_storage {};

struct fortran_storage {};


            namespace detail {

template<class storage>
struct ndarray_iterator_traits;

template<>
struct ndarray_iterator_traits<c_storage> {

    static const npy_uint32     index_flag = NPY_ITER_C_INDEX;

    static const NPY_ORDER      order_flag = NPY_CORDER;

};

template<>
struct ndarray_iterator_traits<fortran_storage> {

    static const npy_uint32     index_flag = NPY_ITER_F_INDEX;

    static const NPY_ORDER      order_flag = NPY_FORTRANORDER;

};

template<class T,class Storage,npy_uint32 rw_flag>
class ArrayElementProxy;

            }

template<class T,class Storage,npy_uint32 rw_flag = NPY_ITER_READWRITE>
class NDArrayIterator :
        public boost::iterator_facade<
                NDArrayIterator<T,Storage,rw_flag>,
                T,
                boost::random_access_traversal_tag,
                detail::ArrayElementProxy<T,Storage,rw_flag>
        >   {

    friend class boost::iterator_core_access;
    friend class detail::ArrayElementProxy<T,Storage,rw_flag>;

    PyArrayObject   *m_parray;
    PyArray_Descr   *dtype;

    NpyIter         *m_it;

    char            **m_data_ptr_arr;

    NpyIter_IterNextFunc
                    *iternext;

    npy_intp        nel,        // Total number of elements
                    cur_ind;    // Current 1D index

    bool            is_past_the_end;



public:

    typedef detail::ArrayElementProxy<T,Storage,rw_flag>    reference;

    typedef ptrdiff_t                               difference_type;

private:

    reference dereference() const {
        return reference(*this);
    }

    bool equal(const NDArrayIterator& rhs) const {
        return cur_ind == rhs.cur_ind && *m_data_ptr_arr == *rhs.m_data_ptr_arr;
    }

    void increment() {
        iternext(m_it);
        ++cur_ind;
    }


    void reposition_iterator(npy_intp ind) {
        int         retval(NpyIter_GotoIndex(m_it,ind));
        if (retval == NPY_FAIL) {
            throw std::runtime_error("Could not reposition the iterator");
        }
        cur_ind = ind;
    }

    void decrement() {
        reposition_iterator(cur_ind-1);
    }

    void advance(difference_type n) {
        reposition_iterator(cur_ind+n);
    }

    difference_type distance_to(const NDArrayIterator& rhs) const {
        return rhs.cur_ind - cur_ind;
    }


public:

    NDArrayIterator() : m_parray(NULL),dtype(NULL),m_it(NULL),m_data_ptr_arr(NULL),
            iternext(NULL),nel(0) {}

    NDArrayIterator(PyArrayObject *parray) :
            m_parray(parray),
            dtype(PyArray_DescrFromType(numpy_dtype_traits<T>::typenum)),
            m_it(NpyIter_New(
                    parray,
                    detail::ndarray_iterator_traits<Storage>::index_flag |
                        NPY_ITER_UPDATEIFCOPY | rw_flag,
                    detail::ndarray_iterator_traits<Storage>::order_flag,
                    NPY_UNSAFE_CASTING,
                    dtype)),
            m_data_ptr_arr(NULL),
            iternext(NULL),
            nel(0),cur_ind(0),is_past_the_end(false) {

        if (m_it == NULL) {
            throw std::runtime_error("Could not allocate NDArrayIterator object");
        }

        if (PyArray_INCREF(m_parray)) {
            throw std::runtime_error("Could not incref the NumPy array object.");
        }

        if ((m_data_ptr_arr = NpyIter_GetDataPtrArray(m_it)) == NULL) {
            throw std::runtime_error("Could not get the data array from the NDArrayIterator");
        }

        if ((iternext = NpyIter_GetIterNext(m_it, NULL)) == NULL) {
            throw std::runtime_error("Could not get the iteration function");
        }
        nel = NpyIter_GetIterSize(m_it);
    }

    ~NDArrayIterator() {
        if (NpyIter_Deallocate(m_it) == NPY_FAIL) {
            throw std::runtime_error("Could not deallocate Numpy iterator.");
        }
        m_it = NULL;
        if (PyArray_XDECREF(m_parray)) {
            throw std::runtime_error("Could not properly decref the NumPy array object.");
        }
        Py_XDECREF(dtype);
    }

    NDArrayIterator(const NDArrayIterator& rhs,bool end = false) :
        m_parray(rhs.m_parray),
        dtype(rhs.dtype),
        m_it(NpyIter_Copy(rhs.m_it)),m_data_ptr_arr(NULL),
        iternext(NULL),nel(0),cur_ind(rhs.cur_ind),
        is_past_the_end(end) {

        if (m_it == NULL) {
            throw std::runtime_error("Could not duplicate NumPy array iterator");
        }
        if (PyArray_INCREF(m_parray)) {
            throw std::runtime_error("Could not incref the NumPy array object.");
        }
        Py_XINCREF(dtype);

        if ((m_data_ptr_arr = NpyIter_GetDataPtrArray(m_it)) == NULL) {
            throw std::runtime_error("Could not get the data array from the NDArrayIterator");
        }

        if ((iternext = NpyIter_GetIterNext(m_it, NULL)) == NULL) {
            throw std::runtime_error("Could not get the iteration function");
        }
        nel = NpyIter_GetIterSize(m_it);

        if (is_past_the_end) {
            reposition_iterator(nel-1);
            increment();
        }
    }


    NDArrayIterator& operator =(const NDArrayIterator& rhs) {
        dtype = rhs.dtype; Py_XINCREF(dtype);
        cur_ind = rhs.cur_ind;
        is_past_the_end = rhs.is_past_the_end;

        if ((m_it = NpyIter_Copy(rhs.m_it)) == NULL) {
            throw std::runtime_error("Could not duplicate NumPy array iterator");
        }
        if ((m_data_ptr_arr = NpyIter_GetDataPtrArray(m_it)) == NULL) {
            throw std::runtime_error("Could not get the data array from the NDArrayIterator");
        }

        if ((iternext = NpyIter_GetIterNext(m_it, NULL)) == NULL) {
            throw std::runtime_error("Could not get the iteration function");
        }
        nel = NpyIter_GetIterSize(m_it);

        if (is_past_the_end) {
            reposition_iterator(nel-1);
            increment();
        }
    }


};


            namespace detail {

template<class T,class Storage,npy_uint32 rw_flag>
class ArrayElementProxy {

    typedef NDArrayIterator<T,Storage,rw_flag>  ArrIt;

    ArrIt&      m_arr_it;

    npy_intp    m_ind;          // Index of the element to be retrieved or set

public:

    ArrayElementProxy(const ArrIt& it) :
        m_arr_it(*const_cast<ArrIt*>(&it)),m_ind(it.cur_ind) {}

    ArrayElementProxy& operator =(const T& val) {
        *reinterpret_cast<T*>(*(m_arr_it.m_data_ptr_arr)) = val;
        return *this;
    }

    operator T () const {
        return *reinterpret_cast<T*>(*(m_arr_it.m_data_ptr_arr));
    }
};


template<class T,class Storage>
class ArrayElementProxy<T,Storage,NPY_ITER_READONLY> {

    typedef NDArrayIterator<T,Storage,NPY_ITER_READONLY>    ArrIt;

    ArrIt&      m_arr_it;

    npy_intp    m_ind;          // Index of the element to be retrieved or set

public:

    ArrayElementProxy(const ArrIt& it) :
        m_arr_it(*const_cast<ArrIt*>(&it)),m_ind(it.cur_ind) {}

    operator T () const {
        return *reinterpret_cast<T*>(*(m_arr_it.m_data_ptr_arr));
    }
};


template<class T,class Storage>
class ArrayElementProxy<T,Storage,NPY_ITER_WRITEONLY> {

    typedef NDArrayIterator<T,Storage,NPY_ITER_WRITEONLY>   ArrIt;

    ArrIt&      m_arr_it;

    npy_intp    m_ind;          // Index of the element to be retrieved or set

public:

    ArrayElementProxy(const ArrIt& it) :
        m_arr_it(*const_cast<ArrIt*>(&it)),m_ind(it.cur_ind) {}

    ArrayElementProxy& operator =(const T& val) {
        *reinterpret_cast<T*>(*(m_arr_it.m_data_ptr_arr)) = val;
        return *this;
    }
};

// Partial specializations for complex numbers

template<class T,class Storage>
class ArrayElementProxy<std::complex<T>,Storage,NPY_ITER_READWRITE> {

    static_assert(std::is_floating_point<T>::value,"Non-floating-point complex numbers are not allowed.");

    typedef NDArrayIterator<std::complex<T>,Storage,NPY_ITER_READWRITE> ArrIt;

    typedef typename numpy_dtype_traits< std::complex<T> >::type
                                                        npy_complex_T;

    ArrIt&      m_arr_it;

    npy_intp    m_ind;          // Index of the element to be retrieved or set

public:

    ArrayElementProxy(const ArrIt& it) :
        m_arr_it(*const_cast<ArrIt*>(&it)),m_ind(it.cur_ind) {}

    ArrayElementProxy<std::complex<T>,Storage,NPY_ITER_READWRITE>&
        operator =(const std::complex<T>& val) {

        npy_complex_T   *pel(reinterpret_cast<npy_complex_T*>(*(m_arr_it.m_data_ptr_arr)));

        pel->real = std::real(val);
        pel->imag = std::imag(val);

        return *this;
    }


    operator std::complex<T> () const {

        npy_complex_T   *pel(reinterpret_cast<npy_complex_T*>(*(m_arr_it.m_data_ptr_arr)));

        return std::complex<T>(pel->real,pel->imag);
    }
};


template<class T,class Storage>
class ArrayElementProxy<std::complex<T>,Storage,NPY_ITER_READONLY> {

    static_assert(std::is_floating_point<T>::value,"Non-floating-point complex numbers are not allowed.");

    typedef NDArrayIterator<std::complex<T>,Storage,NPY_ITER_READONLY>  ArrIt;

    typedef typename numpy_dtype_traits< std::complex<T> >::type
                                                        npy_complex_T;

    ArrIt&      m_arr_it;

    npy_intp    m_ind;          // Index of the element to be retrieved or set

public:

    ArrayElementProxy(const ArrIt& it) :
        m_arr_it(*const_cast<ArrIt*>(&it)),m_ind(it.cur_ind) {}

    operator std::complex<T> () const {

        npy_complex_T   *pel(reinterpret_cast<npy_complex_T*>(*(m_arr_it.m_data_ptr_arr)));

        return std::complex<T>(pel->real,pel->imag);
    }
};


template<class T,class Storage>
class ArrayElementProxy<std::complex<T>,Storage,NPY_ITER_WRITEONLY> {

    static_assert(std::is_floating_point<T>::value,"Non-floating-point complex numbers are not allowed.");

    typedef NDArrayIterator<std::complex<T>,Storage,NPY_ITER_WRITEONLY> ArrIt;

    typedef typename numpy_dtype_traits< std::complex<T> >::type
                                                        npy_complex_T;

    ArrIt&      m_arr_it;

    npy_intp    m_ind;          // Index of the element to be retrieved or set

public:

    ArrayElementProxy(const ArrIt& it) :
        m_arr_it(*const_cast<ArrIt*>(&it)),m_ind(it.cur_ind) {}

    ArrayElementProxy<std::complex<T>,Storage,NPY_ITER_WRITEONLY>&
        operator =(const std::complex<T>& val) {

        npy_complex_T   *pel(reinterpret_cast<npy_complex_T*>(*(m_arr_it.m_data_ptr_arr)));

        pel->real = std::real(val);
        pel->imag = std::imag(val);

        return *this;
    }
};


            } // namespace detail


        }
    }
}



#endif /* MDS_UTILS_PYTHON_NUMPY_ARRAY_ITERATOR_HPP_INCLUDED */