loop.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef pRC_TENSOR_TRAIN_COMMON_FUNCTIONS_LOOP_H
#define pRC_TENSOR_TRAIN_COMMON_FUNCTIONS_LOOP_H
 
#include <prc/core/functors/loop.hpp>
#include <prc/core/functors/view.hpp>
#include <prc/tensor_train/operator/views/loop.hpp>
#include <prc/tensor_train/tensor/views/loop.hpp>
 
namespace pRC
{
    template<class F, class... Xs,
        If<All<Any<TensorTrain::IsTensorish<RemoveReference<Xs>>,
            TensorTrain::IsOperatorish<RemoveReference<Xs>>>...>> = 0,
        If<All<IsInvocable<View, Xs>...>> = 0,
        If<IsSame<typename RemoveReference<Xs>::Dimension...>> = 0,
        If<IsInvocable<F, decltype(declval<Xs>().template core<0>())...>> = 0>
    static inline constexpr auto loop(F &&f, Xs &&...args)
    {
        using Dimension = Common<typename RemoveReference<Xs>::Dimension...>;
        return expand(makeSeries<Index, Dimension{}>(),
            [&f, &args...](auto const... seq)
            {
                auto core = [&f, &args...]<Index N>()
                {
                    return forward<F>(f)(
                        forward<Xs>(args).template core<N>()...);
                };
 
                using T = Common<typename decltype(core.template
                                                   operator()<seq>())::Type...>;
 
                if constexpr(Common<
                                 typename decltype(core.template operator()<
                                                   seq>())::Dimension...>() ==
                    3)
                {
                    using Ranks =
                        Common<decltype(chip<0>(
                                   Sizes<decltype(core.template operator()<
                                                  seq>())::size(0)...>())),
                            decltype(chip<Dimension() - 1>(Sizes<
                                decltype(core.template operator()<seq>())::size(
                                    2)...>()))>;
 
                    using N = pRC::Sizes<
                        decltype(core.template operator()<seq>())::size(1)...>;
 
                    return TensorTrain::TensorViews::Loop<T, N, Ranks, F,
                        RemoveReference<decltype(view(forward<Xs>(args)))>...>(
                        forward<F>(f), view(forward<Xs>(args))...);
                }
                if constexpr(Common<
                                 typename decltype(core.template operator()<
                                                   seq>())::Dimension...>() ==
                    4)
                {
                    using Ranks =
                        Common<decltype(chip<0>(
                                   Sizes<decltype(core.template operator()<
                                                  seq>())::size(0)...>())),
                            decltype(chip<Dimension() - 1>(Sizes<
                                decltype(core.template operator()<seq>())::size(
                                    3)...>()))>;
 
                    using M = pRC::Sizes<
                        decltype(core.template operator()<seq>())::size(1)...>;
 
                    using N = pRC::Sizes<
                        decltype(core.template operator()<seq>())::size(2)...>;
 
                    return TensorTrain::OperatorViews::Loop<T, M, N, Ranks, F,
                        RemoveReference<decltype(view(forward<Xs>(args)))>...>(
                        forward<F>(f), view(forward<Xs>(args))...);
                }
            });
    }
 
    template<class F, class... Xs,
        If<All<Any<TensorTrain::IsTensorish<RemoveReference<Xs>>,
            TensorTrain::IsOperatorish<RemoveReference<Xs>>>...>> = 0,
        If<All<IsInvocable<View, Xs>...>> = 0,
        If<IsSame<typename RemoveReference<Xs>::Dimension...>> = 0,
        If<IsInvocable<F, decltype(declval<Xs>().template core<0>())...>> = 0>
    static inline constexpr auto loop(Xs &&...args)
    {
        return loop(F(), forward<Xs>(args)...);
    }
 
    template<class F, class... Xs,
        If<All<Any<TensorTrain::IsTensorish<RemoveReference<Xs>>,
            TensorTrain::IsOperatorish<RemoveReference<Xs>>>...>> = 0,
        If<Not<All<IsInvocable<View, Xs>...>>> = 0,
        If<IsInvocable<Loop<F>, F, Xs &...>> = 0>
    static inline constexpr auto loop(F &&f, Xs &&...args)
    {
        return eval(loop(forward<F>(f), args...));
    }
 
    template<class F, class... Xs,
        If<All<Any<TensorTrain::IsTensorish<RemoveReference<Xs>>,
            TensorTrain::IsOperatorish<RemoveReference<Xs>>>...>> = 0,
        If<Not<All<IsInvocable<View, Xs>...>>> = 0,
        If<IsInvocable<Loop<F>, Xs &...>> = 0>
    static inline constexpr auto loop(Xs &&...args)
    {
        return eval(loop<F>(args...));
    }
}
#endif // pRC_TENSOR_TRAIN_COMMON_FUNCTIONS_LOOP_H