round.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef pRC_TENSOR_TRAIN_COMMON_FUNCTIONS_ROUND_H
#define pRC_TENSOR_TRAIN_COMMON_FUNCTIONS_ROUND_H
 
#include <prc/core/basic/limits.hpp>
#include <prc/core/functors/view.hpp>
#include <prc/core/tensor/functions/reshape.hpp>
#include <prc/core/value/type_traits.hpp>
#include <prc/tensor_train/common/functions/from_cores.hpp>
#include <prc/tensor_train/operator/type_traits.hpp>
#include <prc/tensor_train/tensor/type_traits.hpp>
 
namespace pRC
{
    template<class Ranks, class X, class R = RemoveReference<X>,
        class V = typename R::Value, class VT = V,
        If<All<IsFloat<V>, IsFloat<VT>>> = 0,
        If<Any<TensorTrain::IsTensorish<R>, TensorTrain::IsOperatorish<R>>> = 0,
        If<IsSizes<Ranks>> = 0,
        If<IsSatisfied<(
            typename Ranks::Dimension() == typename R::Dimension() - 1)>> = 0>
    static inline constexpr auto round(X &&a,
        VT const &tolerance = NumericLimits<VT>::epsilon())
    {
        return expand(makeSeries<Index, typename R::Dimension{}>(),
            [&a, &tolerance](auto const... seq)
            {
                auto totalNorm = identity<typename R::Type>();
 
                return RecursiveLambda([&tolerance, &
                    totalNorm]<Direction D = Direction::LeftToRight,
                    Index C = 0, class XA, class XB, class... Xs>(
                    auto const &self, XA &&a, XB &&b, Xs &&...cores) {
                    using RA = RemoveReference<XA>;
                    using RB = RemoveReference<XB>;
 
                    constexpr auto RAL = RA::size(0);
                    constexpr auto RAR = RA::size(typename RA::Dimension() - 1);
 
                    if constexpr(D == Direction::LeftToRight)
                    {
                        if constexpr(C == typename R::Dimension() - 1)
                        {
                            return expand(
                                reverse(makeSeries<Index,
                                    typename R::Dimension{}>()),
                                [&self, &tolerance, &totalNorm](
                                    auto const cores, auto const... seq)
                                {
                                    return self.template
                                    operator()<Direction::RightToLeft, C>(
                                        get<seq>(cores)...);
                                },
                                forwardAsTuple(forward<XB>(b),
                                    forward<Xs>(cores)..., forward<XA>(a)));
                        }
                        else
                        {
                            auto const [q, r] = qr(
                                reshape<RA::size() / RAR, RAR>(forward<XA>(a)));
 
                            auto const newLeft = expand(
                                makeRange<Index, 1,
                                    typename RA::Dimension() - 1>(),
                                [&q = q](auto const... seq)
                                {
                                    return reshape<RAL, RA::size(seq)...,
                                        RemoveReference<decltype(q)>::size(1)>(
                                        q);
                                });
 
                            Tensor rNorm = norm(r);
                            using RN = decltype(rNorm);
 
                            auto const isCloseToZero = rNorm <=
                                identity<RN>(
                                    NumericLimits<typename RN::Value>::min());
 
                            auto const newRight = contract<1, 0>(isCloseToZero
                                    ? eval(zero<decltype(r)>())
                                    : eval(r / rNorm),
                                forward<XB>(b));
 
                            if(isCloseToZero)
                            {
                                totalNorm = zero();
                            }
                            else
                            {
                                totalNorm *= exp(log(rNorm) /
                                    identity<RN>(typename R::Dimension()))();
                            }
 
                            return self.template operator()<D, C + 1>(newRight,
                                forward<Xs>(cores)..., newLeft);
                        }
                    }
                    else
                    {
                        if constexpr(C == 0)
                        {
                            Tensor const firstNorm = norm(b);
                            using RF = decltype(firstNorm);
 
                            auto const isCloseToZero = firstNorm <=
                                identity<RF>(
                                    NumericLimits<typename RF::Value>::min());
 
                            auto const newFirst = isCloseToZero
                                ? eval(zero<RemoveReference<XB>>())
                                : eval(forward<XB>(b) / firstNorm);
 
                            if(isCloseToZero)
                            {
                                totalNorm = zero();
                            }
                            else
                            {
                                totalNorm *= exp(log(firstNorm) /
                                    identity<RF>(typename R::Dimension()))();
                            }
 
                            return expand(
                                reverse(makeSeries<Index,
                                    typename R::Dimension{}>()),
                                [&totalNorm](auto const cores,
                                    auto const... seq)
                                {
                                    return TensorTrain::fromCores(
                                        get<seq>(cores) * totalNorm...);
                                },
                                forwardAsTuple(newFirst, forward<Xs>(cores)...,
                                    forward<XA>(a)));
                        }
                        else
                        {
                            constexpr auto cutoff =
                                min(Ranks::size(C - 1), RAL, RA::size() / RAL);
 
                            auto const [u, s, v] = svd<cutoff>(
                                reshape<RAL, RA::size() / RAL>(forward<XA>(a)),
                                tolerance);
 
                            auto const newRight =
                                expand(makeRange<Index, 1,
                                           typename RA::Dimension() - 1>(),
                                    [&v = v](auto const... seq)
                                    {
                                        return reshape<cutoff, RA::size(seq)...,
                                            RAR>(adjoint(v));
                                    });
 
                            Tensor const newLeft =
                                contract<typename RB::Dimension() - 1, 0>(
                                    forward<XB>(b), eval(u * fromDiagonal(s)));
 
                            return self.template operator()<D, C - 1>(newLeft,
                                forward<Xs>(cores)..., newRight);
                        }
                    }
                })(forward<X>(a).template core<seq>()...);
            });
    }
 
    template<class X, class R = RemoveReference<X>, class V = typename R::Value,
        class VT = V, If<All<IsFloat<V>, IsFloat<VT>>> = 0,
        If<Any<TensorTrain::IsTensorish<R>, TensorTrain::IsOperatorish<R>>> = 0>
    static inline constexpr auto round(X &&a,
        VT const &tolerance = NumericLimits<VT>::epsilon())
    {
        return round<typename R::Ranks>(forward<X>(a), tolerance);
    }
 
    template<Size C, class X, class R = RemoveReference<X>,
        class V = typename R::Value, class VT = V,
        If<All<IsFloat<V>, IsFloat<VT>>> = 0,
        If<Any<TensorTrain::IsTensorish<R>, TensorTrain::IsOperatorish<R>>> = 0>
    static inline constexpr auto round(X &&a,
        VT const &tolerance = NumericLimits<VT>::epsilon())
    {
        return round<decltype(Sizes(
            makeConstantSequence<Index, typename R::Dimension() - 1, C>()))>(
            forward<X>(a), tolerance);
    }
 
    template<Size... Cs, class X, class R = RemoveReference<X>,
        class V = typename R::Value, class VT = V,
        If<All<IsFloat<V>, IsFloat<VT>>> = 0,
        If<Any<TensorTrain::IsTensorish<R>, TensorTrain::IsOperatorish<R>>> = 0,
        If<IsSatisfied<(sizeof...(Cs) > 1)>> = 0,
        If<IsSatisfied<(sizeof...(Cs) == typename R::Dimension() - 1)>> = 0>
    static inline constexpr auto round(X &&a,
        VT const &tolerance = NumericLimits<VT>::epsilon())
    {
        return round<Sizes<Cs...>>(forward<X>(a), tolerance);
    }
}
#endif // pRC_TENSOR_TRAIN_COMMON_FUNCTIONS_ROUND_H