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/tensor/declarations.hpp>
#include <prc/tensor_train/common/functions/from_cores.hpp>
#include <prc/tensor_train/operator/declarations.hpp>
#include <prc/tensor_train/tensor/declarations.hpp>
 
namespace pRC
{
    template<IsSizes Ranks, class X, class R = RemoveReference<X>,
        IsFloat VT = Value<R>>
        requires(TensorTrain::IsTensorish<R> ||
                    TensorTrain::IsOperatorish<R>) &&
        IsFloat<Value<R>> && (Ranks::Dimension == R::Dimension - 1)
    static inline constexpr auto round(X &&a,
        VT const &tolerance = NumericLimits<VT>::epsilon())
    {
        return expand(makeSeries<Index, 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(RA::Dimension - 1);
 
                        if constexpr(D == Direction::LeftToRight)
                        {
                            if constexpr(C == R::Dimension - 1)
                            {
                                return expand(
                                    reverse(makeSeries<Index, 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, 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<Value<RN>>::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>(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<Value<RF>>::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>(R::Dimension))();
                                }
 
                                return expand(
                                    reverse(makeSeries<Index, 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, RA::Dimension - 1>(),
                                    [&v = v](auto const... seq)
                                    {
                                        return reshape<cutoff, RA::size(seq)...,
                                            RAR>(adjoint(v));
                                    });
 
                                Tensor const newLeft =
                                    contract<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>, IsFloat VT = Value<R>>
        requires(TensorTrain::IsTensorish<R> ||
                    TensorTrain::IsOperatorish<R>) &&
        IsFloat<Value<R>>
    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>,
        IsFloat VT = Value<R>>
        requires(TensorTrain::IsTensorish<R> ||
                    TensorTrain::IsOperatorish<R>) &&
        IsFloat<Value<R>>
    static inline constexpr auto round(X &&a,
        VT const &tolerance = NumericLimits<VT>::epsilon())
    {
        return round<
            decltype(makeConstantSequence<Size, R::Dimension - 1, C>())>(
            forward<X>(a), tolerance);
    }
 
    template<Size... Cs, class X, class R = RemoveReference<X>,
        IsFloat VT = Value<R>>
        requires(TensorTrain::IsTensorish<R> ||
                    TensorTrain::IsOperatorish<R>) &&
        IsFloat<Value<R>> && (sizeof...(Cs) > 1) &&
        (sizeof...(Cs) == R::Dimension - 1)
    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