ndim.hpp

/****************************************************************************
 *  hipipe library
 *  Copyright (c) 2017, Cognexa Solutions s.r.o.
 *  Copyright (c) 2018, Iterait a.s.
 *  Author(s) Filip Matzner
 *
 *  This file is distributed under the MIT License.
 *  See the accompanying file LICENSE.txt for the complete license agreement.
 ****************************************************************************/
 
#pragma once
 
#include <hipipe/core/utility/random.hpp>
 
#include <range/v3/action/reverse.hpp>
#include <range/v3/algorithm/adjacent_find.hpp>
#include <range/v3/algorithm/all_of.hpp>
#include <range/v3/algorithm/fill.hpp>
#include <range/v3/algorithm/find.hpp>
#include <range/v3/algorithm/max.hpp>
#include <range/v3/numeric/accumulate.hpp>
#include <range/v3/numeric/partial_sum.hpp>
#include <range/v3/view/all.hpp>
#include <range/v3/view/chunk.hpp>
#include <range/v3/view/for_each.hpp>
 
#include <cassert>
#include <cmath>
#include <functional>
#include <memory>
#include <random>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace hipipe::utility {
 
namespace rga = ranges::actions;
namespace rgv = ranges::views;
 
template<typename Rng, typename PrevRng = void, bool IsRange = ranges::range<Rng>>
struct ndims {
};
 
template<typename Rng, typename PrevRng>
struct ndims<Rng, PrevRng, false> : std::integral_constant<long, 0L> {
};
 
template<typename Rng, typename PrevRng>
struct ndims<Rng, PrevRng, true>
  : std::integral_constant<long, ndims<ranges::range_value_t<Rng>, Rng>{} + 1> {
};
 
// For recursive ranges, such as std::filesystem::path, do not recurse further and
// consider it to be a scalar.
template<typename Rng>
struct ndims<Rng, Rng, true> : std::integral_constant<long, -1L> {
};
 
template<typename, template<typename...> typename>
struct is_specialization : std::false_type {};
 
template<template<typename...> typename Template, typename... Args>
struct is_specialization<Template<Args...>, Template> : std::true_type {};
 
template<typename T, long Dims, template<typename...> typename Container = std::vector>
struct ndim_container {
    using type = Container<typename ndim_container<T, Dims-1, Container>::type>;
    static_assert(Dims >= 0, "The number of dimensions has to be non-negative.");
};
 
template<typename T, template<typename...> typename Container>
struct ndim_container<T, 0, Container> {
    using type = T;
};
 
template<typename T, long Dims, template<typename...> typename Container = std::vector>
using ndim_container_t = typename ndim_container<T, Dims, Container>::type;
 
template<typename Rng, long Dim = -1L>
struct ndim_type
  : ndim_type<typename ranges::range_value_t<Rng>, Dim - 1L> {
  static_assert(Dim > 0, "Dimension has to be positive, zero or -1.");
};
 
template<typename T>
struct ndim_type<T, 0L> {
    using type = T;
};
 
template<typename Rng>
struct ndim_type<Rng, -1L>
  : ndim_type<Rng, ndims<Rng>{}> {
};
 
template<typename T, long Dim = -1L>
using ndim_type_t = typename ndim_type<T, Dim>::type;
 
// multidimensional range size //
 
namespace detail {
 
    template<typename Rng, long Dim, long NDims>
    struct ndim_size_impl {
        static void impl(const Rng& rng, std::vector<std::vector<long>>& size_out)
        {
            size_out[Dim-1].push_back(ranges::size(rng));
            for (auto& subrng : rng) {
                ndim_size_impl<ranges::range_value_t<Rng>, Dim+1, NDims>
                  ::impl(subrng, size_out);
            }
        }
    };
 
    template<typename Rng, long Dim>
    struct ndim_size_impl<Rng, Dim, Dim> {
        static void impl(const Rng& rng, std::vector<std::vector<long>>& size_out)
        {
            size_out[Dim-1].push_back(ranges::size(rng));
        }
    };
 
}  // namespace detail
 
template<long NDims, typename Rng>
std::vector<std::vector<long>> ndim_size(const Rng& rng)
{
    static_assert(NDims > 0);
    std::vector<std::vector<long>> size_out(NDims);
    detail::ndim_size_impl<Rng, 1, NDims>::impl(rng, size_out);
    return size_out;
}
 
template<typename Rng>
std::vector<std::vector<long>> ndim_size(const Rng& rng)
{
    return utility::ndim_size<ndims<Rng>{}>(rng);
}
 
// multidimensional range resize //
 
namespace detail {
 
    template<long Dim, long NDims>
    struct ndim_resize_impl {
        template<typename Rng, typename ValT>
        static void impl(Rng& vec,
                         const std::vector<std::vector<long>>& vec_size,
                         std::vector<long>& vec_size_idx,
                         const ValT& val)
        {
            vec.resize(vec_size[Dim-1][vec_size_idx[Dim-1]++]);
            for (auto& subvec : vec) {
                ndim_resize_impl<Dim+1, NDims>::impl(subvec, vec_size, vec_size_idx, val);
            }
        }
    };
 
    template<long Dim>
    struct ndim_resize_impl<Dim, Dim> {
        template<typename Rng, typename ValT>
        static void impl(Rng& vec,
                         const std::vector<std::vector<long>>& vec_size,
                         std::vector<long>& vec_size_idx,
                         const ValT& val)
        {
            vec.resize(vec_size[Dim-1][vec_size_idx[Dim-1]++], val);
        }
    };
 
}  // namespace detail
 
template<long NDims, typename Rng, typename ValT = ndim_type_t<Rng, NDims>>
Rng& ndim_resize(Rng& vec,
                 const std::vector<std::vector<long>>& vec_size,
                 ValT val = ValT{})
{
    // check that the size is valid
    assert(vec_size.size() == NDims);
    static_assert(NDims <= ndims<Rng>{} - ndims<ValT>{});
    for (std::size_t i = 1; i < vec_size.size(); ++i) {
        assert(vec_size[i].size() == ranges::accumulate(vec_size[i-1], 0UL));
    }
    // build initial indices
    std::vector<long> vec_size_idx(vec_size.size());
    // recursively resize
    detail::ndim_resize_impl<1, NDims>::impl(vec, vec_size, vec_size_idx, val);
    return vec;
}
 
template<typename Rng, typename ValT = ndim_type_t<Rng>>
Rng& ndim_resize(Rng& vec,
                 const std::vector<std::vector<long>>& vec_size,
                 ValT val = ValT{})
{
    return ndim_resize<ndims<Rng>{}-ndims<ValT>{}>(vec, vec_size, std::move(val));
}
 
template<long NDims, typename Rng, typename ValT = ndim_type_t<Rng, NDims>>
Rng& ndim_pad(Rng& vec, ValT val = ValT{})
{
    static_assert(NDims <= ndims<Rng>{} - ndims<ValT>{});
    std::vector<std::vector<long>> vec_size = ndim_size<NDims>(vec);
    // replace the sizes in each dimension with the max size in the same dimension
    for (std::size_t i = 1; i < vec_size.size(); ++i) {
        long max_size = ranges::max(vec_size[i]);
        ranges::fill(vec_size[i], max_size);
    }
    return utility::ndim_resize<NDims>(vec, vec_size, std::move(val));
}
 
template<typename Rng, typename ValT = ndim_type_t<Rng>>
Rng& ndim_pad(Rng& vec, ValT val = ValT{})
{
    return utility::ndim_pad<ndims<Rng>{}-ndims<ValT>{}>(vec, std::move(val));
}
 
// multidimensional range shape //
 
namespace detail {
 
    template<typename Rng, long Dim, long NDims>
    struct shape_impl {
        static void impl(const Rng& rng, std::vector<long>& shape)
        {
            shape[Dim-1] = ranges::size(rng);
            if (ranges::size(rng)) {
                shape_impl<typename ranges::range_value_t<Rng>, Dim+1, NDims>
                  ::impl(*ranges::begin(rng), shape);
            }
        }
    };
 
    template<typename Rng, long Dim>
    struct shape_impl<Rng, Dim, Dim> {
        static void impl(const Rng& rng, std::vector<long>& shape)
        {
            shape[Dim-1] = ranges::size(rng);
        }
    };
 
    // this is just for assertion purposes - check that the entire vector has a rectangular shape
    template<typename Rng, long NDims>
    struct check_rectangular_shape {
        static bool all_same(std::vector<long>& vec)
        {
            return ranges::adjacent_find(vec, std::not_equal_to<long>{}) == ranges::end(vec);
        }
 
        static bool impl(const Rng& rng)
        {
            std::vector<std::vector<long>> size = utility::ndim_size<NDims>(rng);
            return ranges::all_of(size, all_same);
        }
    };
 
}  // namespace detail
 
template<long NDims, typename Rng>
std::vector<long> shape(const Rng& rng)
{
    static_assert(NDims > 0);
    std::vector<long> shape(NDims);
    // the ndim_size is not used for efficiency in ndebug mode (only the 0-th element is inspected)
    detail::shape_impl<Rng, 1, NDims>::impl(rng, shape);
    assert((detail::check_rectangular_shape<Rng, NDims>::impl(rng)));
    return shape;
}
 
template<typename Rng>
std::vector<long> shape(const Rng& rng)
{
    return utility::shape<ndims<Rng>{}>(rng);
}
 
// recursive range flatten //
 
namespace detail {
 
    // recursively join the vector
    template<long Dim>
    struct flat_view_impl {
        static auto impl()
        {
            return rgv::for_each([](auto& subrng) {
                return flat_view_impl<Dim-1>::impl()(subrng);
            });
        }
    };
 
    // for 0, return the original vector
    template<>
    struct flat_view_impl<1> {
        static auto impl()
        {
            return rgv::all;
        }
    };
 
}  // namespace detail
 
template<long NDims, typename Rng>
auto flat_view(Rng& rng)
{
    static_assert(NDims > 0);
    return detail::flat_view_impl<NDims>::impl()(rng);
}
 
template<long NDims, typename Rng>
auto flat_view(const Rng& rng)
{
    static_assert(NDims > 0);
    return detail::flat_view_impl<NDims>::impl()(rng);
}
 
template<typename Rng>
auto flat_view(Rng&& rng)
{
    return utility::flat_view<ndims<Rng>{}>(std::forward<Rng>(rng));
}
 
// reshaped view of a multidimensional range //
 
namespace detail {
 
    template<long N>
    struct reshaped_view_impl_go {
        static auto impl(const std::shared_ptr<std::vector<long>>& shape_ptr)
        {
            return rgv::chunk((*shape_ptr)[N-2])
              | rgv::transform([shape_ptr](auto subview) {
                    return reshaped_view_impl_go<N-1>::impl(shape_ptr)(std::move(subview));
            });
        }
    };
 
    template<>
    struct reshaped_view_impl_go<1> {
        static auto impl(const std::shared_ptr<std::vector<long>>&)
        {
            return rgv::all;
        }
    };
 
    template<long N, typename Rng>
    auto reshaped_view_impl(Rng& vec, std::vector<long> shape)
    {
        assert(shape.size() == N);
        auto flat = flat_view(vec);
 
        // if -1 present in the shape list, deduce the dimension
        auto deduced_pos = ranges::find(shape, -1);
        if (deduced_pos != shape.end()) {
            auto flat_size = ranges::distance(flat);
            auto shape_prod = -ranges::accumulate(shape, 1, std::multiplies<>{});
            assert(flat_size % shape_prod == 0);
            *deduced_pos = flat_size / shape_prod;
        }
 
        // check that all the requested dimenstions have positive size
        assert(ranges::all_of(shape, [](long s) { return s > 0; }));
        // check that the user requests the same number of elements as there really is
        assert(ranges::distance(flat) == ranges::accumulate(shape, 1, std::multiplies<>{}));
        // calculate the cummulative product of the shape list in reverse order
        shape |= rga::reverse;
        ranges::partial_sum(shape, shape, std::multiplies<>{});
        // the recursive chunks will share a single copy of the shape list (performance)
        auto shape_ptr = std::make_shared<std::vector<long>>(std::move(shape));
        return detail::reshaped_view_impl_go<N>::impl(shape_ptr)(std::move(flat));
    }
 
}  // namespace detail
 
template<long N, typename Rng>
auto reshaped_view(Rng& rng, std::vector<long> shape)
{
    return detail::reshaped_view_impl<N, Rng>(rng, std::move(shape));
}
 
template<long N, typename Rng>
auto reshaped_view(const Rng& rng, std::vector<long> shape)
{
    return detail::reshaped_view_impl<N, const Rng>(rng, std::move(shape));
}
 
// generate //
 
namespace detail {
 
    template<long Dim, long NDims>
    struct generate_impl {
        template<typename Rng, typename Gen>
        static void impl(Rng& rng, Gen& gen, long gendims)
        {
            if (Dim > gendims) {
                auto val = std::invoke(gen);
                for (auto& elem : flat_view<NDims-Dim+1>(rng)) elem = val;
            } else {
                for (auto& subrng : rng) {
                    detail::generate_impl<Dim+1, NDims>::impl(subrng, gen, gendims);
                }
            }
        }
    };
 
    template<long Dim>
    struct generate_impl<Dim, Dim> {
        template<typename Rng, typename Gen>
        static void impl(Rng& rng, Gen& gen, long gendims)
        {
            if (Dim > gendims) ranges::fill(rng, std::invoke(gen));
            else for (auto& val : rng) val = std::invoke(gen);
        }
    };
 
}  // namespace detail
 
template<long NDims, typename Rng, typename Gen>
void generate(Rng&& rng,
              Gen&& gen,
              long gendims = std::numeric_limits<long>::max())
{
    static_assert(NDims > 0);
    detail::generate_impl<1, NDims>::impl(rng, gen, gendims);
}
 
template<typename Rng, typename Gen>
void generate(Rng&& rng,
              Gen&& gen,
              long gendims = std::numeric_limits<long>::max())
{
    using GenT = std::result_of_t<Gen()>;
    detail::generate_impl<1, ndims<Rng>{}-ndims<GenT>{}>::impl(rng, gen, gendims);
}
 
template<long NDims,
         typename Rng,
         typename Prng = std::mt19937&,
         typename Dist = std::uniform_real_distribution<double>>
void random_fill(Rng&& rng,
                 Dist&& dist = Dist{0, 1},
                 Prng&& prng = utility::random_generator,
                 long gendims = std::numeric_limits<long>::max())
{
    auto gen = [&dist, &prng]() { return std::invoke(dist, prng); };
    utility::generate<NDims>(std::forward<Rng>(rng), std::move(gen), gendims);
}
 
template<typename Rng,
         typename Prng = std::mt19937&,
         typename Dist = std::uniform_real_distribution<double>>
void random_fill(Rng&& rng,
                 Dist&& dist = Dist{0, 1},
                 Prng&& prng = utility::random_generator,
                 long gendims = std::numeric_limits<long>::max())
{
    auto gen = [&dist, &prng]() { return std::invoke(dist, prng); };
    utility::generate(std::forward<Rng>(rng), std::move(gen), gendims);
}
 
namespace detail {
 
    struct same_size_impl {
        template<typename Rng, typename... Rngs>
        bool operator()(Rng&& rng, Rngs&&... rngs) const
        {
            return (... && (ranges::size(rng) == ranges::size(rngs)));
        }
 
        bool operator()() const
        {
            return true;
        }
    };
 
}  // namespace detail
 
template<typename Tuple>
bool same_size(Tuple&& rngs)
{
    return std::apply(detail::same_size_impl{}, rngs);
}
 
}  // namespace hipipe::utility