core/grid.hh

#ifndef VTESTBED_CORE_GRID_HH
#define VTESTBED_CORE_GRID_HH

#include <cmath>
#include <ostream>
#include <stdexcept>

#include <vtestbed/base/blitz.hh>
#include <vtestbed/base/bstream.hh>
#include <vtestbed/base/for_loop.hh>

namespace vtb {

    namespace core {

        template <class T, int N>
        class Grid {

        public:
            typedef blitz::TinyVector<T,N> vec;
            typedef blitz::TinyVector<int,N> int_n;

        private:
            vec _lbound{T{0}};
            vec _ubound{T{0}};
            int_n _begin_index{blitz::const_vector<int,N>(0)};
            int_n _end_index{blitz::const_vector<int,N>(0)};

        public:

            Grid() = default;
            ~Grid() = default;
            Grid(const Grid&) = default;
            Grid& operator=(const Grid&) = default;
            Grid(Grid&&) = default;
            Grid& operator=(Grid&&) = default;

            template <class X>
            inline explicit
            Grid(const Grid<X,N>& rhs) {
                this->_lbound = rhs.lbound();
                this->_ubound = rhs.ubound();
                this->_begin_index = rhs.begin();
                this->_end_index = rhs.end();
            }

            inline
            Grid(
                const vec& lbound,
                const vec& ubound,
                const int_n& begin_index,
                const int_n& end_index
            ):
            _lbound(lbound),
            _ubound(ubound),
            _begin_index(begin_index),
            _end_index(end_index) {
                #if defined(VTB_DEBUG)
                for (int i=0; i<N; ++i) {
                    if (this->begin(i) > this->end(i)) {
                        throw std::invalid_argument(
                            "vtb::core::Grid: bad begin/end index"
                        );
                    }
                }
                #endif
            }

            inline
            Grid(const vec& lbound, const vec& ubound, const int_n& shape):
            Grid(lbound, ubound, blitz::const_vector<int, N>(0), shape) {}

            inline
            Grid(const vec& ubound, const int_n& shape):
            Grid(
                blitz::const_vector<T,N>(T{0}),
                ubound,
                blitz::const_vector<int,N>(0),
                shape
            ) {}

            inline explicit
            Grid(const int_n& shape):
            Grid(vec{shape-1}, shape) {}

            template <class ... Tail>
            inline explicit
            Grid(const Grid<T,1>& head, const Tail& ... tail) {
                const Grid<T,1>* grids[N] = {&head, &(tail)...};
                for (int i=0; i<N; ++i) {
                    this->_lbound(i) = grids[i]->lbound(0);
                    this->_ubound(i) = grids[i]->ubound(0);
                    this->_begin_index(i) = grids[i]->begin(0);
                    this->_end_index(i) = grids[i]->end(0);
                }
            }

            inline const int_n
            shape() const noexcept {
                return this->end() - this->begin();
            }

            inline int
            extent(int i) const noexcept {
                return this->end(i) - this->begin(i);
            }

            inline const int_n&
            begin() const noexcept {
                return this->_begin_index;
            }

            inline int_n&
            begin() noexcept {
                return this->_begin_index;
            }

            inline int
            begin(int i) const noexcept {
                return this->_begin_index(i);
            }

            inline int&
            begin(int i) noexcept {
                return this->_begin_index(i);
            }

            inline const int_n&
            end() const noexcept {
                return this->_end_index;
            }

            inline int_n&
            end() noexcept {
                return this->_end_index;
            }

            inline int&
            end(int i) noexcept {
                return this->_end_index(i);
            }

            inline int
            end(int i) const noexcept {
                return this->_end_index(i);
            }

            inline const vec&
            lbound() const noexcept {
                return this->_lbound;
            }

            inline vec&
            lbound() noexcept {
                return this->_lbound;
            }

            inline T
            lbound(int i) const noexcept {
                return this->_lbound(i);
            }

            inline T&
            lbound(int i) noexcept {
                return this->_lbound(i);
            }

            inline const vec&
            ubound() const noexcept {
                return this->_ubound;
            }

            inline vec&
            ubound() noexcept {
                return this->_ubound;
            }

            inline T
            ubound(int i) const noexcept {
                return this->_ubound(i);
            }

            inline T&
            ubound(int i) noexcept {
                return this->_ubound(i);
            }

            inline int
            num_patches(int i) const {
                return std::max(this->extent(i)-1, 1);
            }

            inline const int_n
            num_patches() const {
                return blitz::max(this->shape()-1, 1);
            }

            inline const vec
            size() const noexcept {
                return this->ubound() - this->lbound();
            }

            inline const T
            length(int i) const noexcept {
                return this->ubound(i) - this->lbound(i);
            }

            inline size_t num_elements() const { return product(shape()); }

            inline const T
            delta(int i) const {
                const int npatches = this->num_patches(i);
                return npatches == 0 ? T(0) : (this->length(i) / npatches);
            }

            inline vec
            delta() const {
                using blitz::where;
                const auto& npatches = this->num_patches();
                return where(npatches == 0, T(0), this->size() / npatches);
            }

            inline const vec
            operator()(const int_n& i) const noexcept {
                return this->lbound() + this->delta()*(i - this->begin());
            }

            inline const vec
            operator()(const int_n& i, const vec& delta) const noexcept {
                return this->lbound() + delta*(i - this->begin());
            }

            inline const T
            operator()(const int i, const int dim) const noexcept {
                return this->lbound(dim) +
                    this->delta(dim)*(i - this->begin(dim));
            }

            inline const vec
            index(const vec& x) const noexcept {
                using blitz::where;
                const auto& npatches = this->num_patches();
                return this->begin() + where(
                    npatches == 0,
                    T(0),
                    (x - this->lbound()) / this->delta()
                );
            }

            inline const T
            index(T x, const int dim) const noexcept {
                const int npatches = this->num_patches(dim);
                return this->begin(dim) +
                    (npatches == 0 ? T(0)
                     : ((x - this->lbound(dim)) / this->delta(dim)));
            }

            inline Grid<T,N>
            operator()(const blitz::RectDomain<N>& domain) const {
                Grid<T,N> result;
                for (int i=0; i<N; ++i) {
                    auto start = domain.lbound(i);
                    auto end = domain.ubound(i);
                    result._begin_index[i] = start;
                    result._end_index[i] = end+1;
                    result._lbound[i] = (*this)(start, i);
                    result._ubound[i] = (*this)(end, i);
                }
                return result;
            }

            template <class ... Tail>
            inline Grid<T,N>
            operator()(const blitz::Range& head, const Tail& ... tail) const {
                const blitz::Range* ranges[N] = {&head, &(tail)...};
                Grid<T,N> result;
                for (int i=0; i<N; ++i) {
                    const auto& r = *ranges[i];
                    auto start = r.first(this->_begin_index[i]);
                    auto end = r.last(this->_end_index[i]-1);
                    result._begin_index[i] = start;
                    result._end_index[i] = end+1;
                    result._lbound[i] = (*this)(start, i);
                    result._ubound[i] = (*this)(end, i);
                }
                return result;
            }

            template <class ... Tail>
            inline auto
            select(int head, Tail ... tail) const -> Grid<T,sizeof...(tail)+1> {
                const int dimensions[] = {head, tail...};
                Grid<T,sizeof...(tail)+1> result;
                for (size_t i=0; i<sizeof...(tail)+1; ++i) {
                    const auto dim = dimensions[i];
                    if (dim >= N) { throw std::invalid_argument("bad dimension"); }
                    result.begin(i) = this->_begin_index[dim],
                    result.end(i) = this->_end_index[dim],
                    result.lbound(i) = this->_lbound[dim],
                    result.ubound(i) = this->_ubound[dim];
                }
                return result;
            }

            inline void compact() { end() -= begin(); begin() = 0; }

            inline void
            clear() noexcept {
                this->_lbound = T(0);
                this->_ubound = T(0);
                this->_begin_index = 0;
                this->_end_index = 1;
            }

            inline bool
            within_bounds(T x, int dim) const noexcept {
                return !(x < this->_lbound(dim) || x > this->_ubound(dim));
            }

            inline bool
            contains(const vec& x) const noexcept {
                return !blitz::any(x < this->lbound() || x > this->ubound());
            }

            template <class ... Args>
            inline bool
            contains(const vec& x , const Args& ... tail) const noexcept {
                return this->contains(x) && contains(tail...);
            }

            inline bool
            near(const Grid& rhs, T eps) {
                using blitz::all;
                using blitz::near;
                return
                    all(rhs.begin() == this->begin()) &&
                    all(rhs.end() == this->end()) &&
                    near(rhs.lbound(), this->lbound(), eps) &&
                    near(rhs.ubound(), this->ubound(), eps);
            }

            inline bool
            empty() const {
                using blitz::any;
                return any(lbound() > ubound()) || any(shape() <= 1);
            }

            inline friend bstream&
            operator<<(bstream& out, const Grid& rhs) {
                out << rhs._lbound;
                out << rhs._ubound;
                out << rhs._begin_index;
                out << rhs._end_index;
                return out;
            }

            inline friend bstream&
            operator>>(bstream& in, Grid& rhs) {
                in >> rhs._lbound;
                in >> rhs._ubound;
                in >> rhs._begin_index;
                in >> rhs._end_index;
                return in;
            }

            friend std::ostream&
            operator<<(std::ostream& out, const Grid& rhs) {
                return out
                    << "from" << ' ' << rhs._lbound << ' '
                    << "to" << ' ' << rhs._ubound << ' '
                    << "start" << ' ' << rhs._begin_index << ' '
                    << "end" << ' ' << rhs._end_index;
            }

        };

        template <class T>
        blitz::Array<T,1>
        generate(const Grid<T,1>& grid);

        template <class T, int N>
        inline blitz::TinyVector<T,N>
        clamp(const blitz::TinyVector<T,N> x, const Grid<T,N>& grid) {
            using blitz::min;
            using blitz::max;
            return max(grid.lbound(), min(grid.ubound(), x));
        }

        template <class T, int N, class Function>
        inline void
        generate(
            const Grid<T,N>& grid,
            blitz::Array<T,N>& result,
            Function func
        ) {
            #if defined(VTB_DEBUG)
            if (!all(result.shape() <= grid.end())) {
                throw std::invalid_argument(
                    "vtb::core::generate: bad result shape"
                );
            }
            #endif
            const auto delta{grid.delta()};
            parallel_for_loop<N>(
                grid.begin(),
                grid.end(),
                [&](const blitz::TinyVector<int,N>& idx) {
                    result(idx) = func(grid(idx, delta));
                }
            );
        }

        template <class T, int N, class Function>
        inline blitz::Array<T,N>
        generate(
            const Grid<T,N>& grid,
            Function func
        ) {
            blitz::Array<T,N> result(grid.shape());
            generate<T,N,Function>(grid, result, func);
            return result;
        }

    }

}

#endif // vim:filetype=cpp