unit.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef pRC_TENSOR_TRAIN_TENSOR_UNIT_H
#define pRC_TENSOR_TRAIN_TENSOR_UNIT_H
 
#include <prc/core/basic/functions/unit.hpp>
#include <prc/core/functors/eval.hpp>
#include <prc/core/tensor/functions/direct_sum.hpp>
#include <prc/tensor_train/tensor/views/enumerate.hpp>
 
namespace pRC
{
    template<class T>
    struct Unit<T, If<TensorTrain::IsTensorView<T>>> : Unit<ResultOf<Eval, T>>
    {
    };
 
    template<class T>
    struct Unit<T, If<TensorTrain::IsTensor<T>>>
    {
        constexpr auto operator()()
        {
            auto const f = []<Index C>()
            {
                using Core = typename T::template Cores<C>;
                constexpr auto CRL = Core::size(0);
                constexpr auto CN = Core::size(1);
                constexpr auto CRR = Core::size(2);
 
                return permute<1, 0, 2>(exclude<DirectSum, 1>(
                    unit<typename Core::template ChangeSizes<1, CN, 1>>(),
                    zero<typename Core::template ChangeSizes<CRL - 1, CN,
                        CRR - 1>>()));
            };
 
            using F = RemoveConstReference<decltype(f)>;
            using N = typename T::N;
            using Ranks = typename T::Ranks;
 
            return TensorTrain::TensorViews::Enumerate<typename T::Type, N,
                Ranks, F>(f);
        }
 
        template<class X, If<IsConstructible<typename T::Type, X>> = 0>
        constexpr auto operator()(X &&value)
        {
            auto const f =
                [value = typename T::Type(forward<X>(value))]<Index C>()
            {
                using Core = typename T::template Cores<C>;
                constexpr auto CRL = Core::size(0);
                constexpr auto CN = Core::size(1);
                constexpr auto CRR = Core::size(2);
 
                if constexpr(C == 0)
                {
                    return permute<1, 0, 2>(exclude<DirectSum, 1>(
                        unit<typename Core::template ChangeSizes<1, CN, 1>>(
                            value),
                        zero<typename Core::template ChangeSizes<CRL - 1, CN,
                            CRR - 1>>()));
                }
                else
                {
                    return permute<1, 0, 2>(exclude<DirectSum, 1>(
                        unit<typename Core::template ChangeSizes<1, CN, 1>>(),
                        zero<typename Core::template ChangeSizes<CRL - 1, CN,
                            CRR - 1>>()));
                }
            };
 
            using F = RemoveConstReference<decltype(f)>;
            using N = typename T::N;
            using Ranks = typename T::Ranks;
 
            return TensorTrain::TensorViews::Enumerate<typename T::Type, N,
                Ranks, F>(f);
        }
    };
}
#endif // pRC_TENSOR_TRAIN_TENSOR_UNIT_H