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/direct_sum.hpp>
#include <prc/core/functors/eval.hpp>
#include <prc/tensor_train/tensor/views/enumerate.hpp>
 
namespace pRC
{
    template<TensorTrain::IsTensor T>
    struct Unit<T>
    {
        constexpr auto operator()() const
        {
            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>
            requires IsConstructible<typename T::Type, X>
        constexpr auto operator()(X &&value) const
        {
            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);
        }
    };
 
    template<TensorTrain::IsTensorView T>
    struct Unit<T> : Unit<ResultOf<Eval, T>>
    {
    };
}
#endif // pRC_TENSOR_TRAIN_TENSOR_UNIT_H