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/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>
        requires((TensorTrain::IsTensorish<RemoveReference<Xs>> ||
                     TensorTrain::IsOperatorish<RemoveReference<Xs>>) &&
                    ...) &&
        (sizeof...(Xs) > 0) && (isSame(RemoveReference<Xs>::Dimension...)) &&
        requires { declval<F>()(declval<Xs>().template core<0>()...); }
    static inline constexpr auto loop(F &&f, Xs &&...args)
    {
        if constexpr((IsInvocable<View, Xs> && ...))
        {
            constexpr auto Dimension =
                Common<Constant<Size, 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<Constant<Size,
                                     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<Constant<Size,
                                     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 =
                            Sizes<decltype(core.template
                                           operator()<seq>())::size(1)...>;
 
                        using N =
                            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))...);
                    }
                });
        }
        else
        {
            return eval(loop(forward<F>(f), args...));
        }
    }
 
    template<class F, class... Xs>
        requires((TensorTrain::IsTensorish<RemoveReference<Xs>> ||
                     TensorTrain::IsOperatorish<RemoveReference<Xs>>) &&
                    ...) &&
        requires { loop(declval<F>(), declval<Xs>()...); }
    static inline constexpr auto loop(Xs &&...args)
    {
        return loop(F(), forward<Xs>(args)...);
    }
}
#endif // pRC_TENSOR_TRAIN_COMMON_FUNCTIONS_LOOP_H