ops/lu.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/linalg/backends/cpu/lu.hpp>
#include <uni20/linalg/backends/lapack/lu.hpp>
#if UNI20_HAS_MPLAPACK_FLOAT80
#include <uni20/linalg/backends/lapack/lu_binary80.hpp>
#endif
#include <optional>
#include <span>
#include <uni20/linalg/precision.hpp>
#include <uni20/tensor/copy.hpp>
#include <uni20/tensor/precision.hpp>
#include <uni20/tensor/tensor.hpp>
#include <vector>
 
namespace uni20::linalg
{
namespace detail
{
struct LuFactorAccess;
}
 
template <RealOrComplex S> class LuFactorization {
  public:
    using value_type = S;
    std::size_t order() const noexcept { return std::size_t(factors_.extent(0)); }
    DenseMatrix<S> const& packed() const& noexcept { return factors_; }
    std::span<std::size_t const> pivots() const& noexcept { return pivots_; }
    Precision precision() const { return precision_of(factors_); }
    std::optional<LuNorms<make_real_t<S>>> const& norms() const& noexcept { return norms_; }
 
  private:
    friend struct detail::LuFactorAccess;
    LuFactorization(DenseMatrix<S>&& factors, std::vector<std::size_t>&& pivots,
                    std::optional<LuNorms<make_real_t<S>>>&& norms)
        : factors_(std::move(factors)), pivots_(std::move(pivots)), norms_(std::move(norms))
    {}
    DenseMatrix<S> factors_;
    std::vector<std::size_t> pivots_;
    std::optional<LuNorms<make_real_t<S>>> norms_;
};
 
template <RealOrComplex S> struct LuFactorizationResult
{
    SolveInfo info;
    std::optional<LuFactorization<S>> factor;
};
 
namespace detail
{
struct LuFactorAccess
{
    template <class S>
    static auto make(DenseMatrix<S>&& factors, std::vector<std::size_t>&& pivots,
                     std::optional<LuNorms<make_real_t<S>>>&& norms)
    {
      return LuFactorization<S>(std::move(factors), std::move(pivots), std::move(norms));
    }
};
 
template <KernelBackendSelector Backend, RankedTensorView<2> A>
auto lu_factor_impl(Backend&& backend, A const& a, LuOptions<make_real_t<tensor_element_t<A>>> const& options,
                    Precision p)
{
  using S = tensor_element_t<A>;
  ERROR_IF(a.extent(0) != a.extent(1), "LU requires a square matrix");
  SolveOptions<make_real_t<S>> pivot_options{options.relative_pivot_tolerance};
  detail::require_solve_options(pivot_options);
  require_precision<S>(p);
  (void)p.bit_count();
  auto work = make_tensor<ColumnMajor>(a);
  std::vector<std::size_t> pivots(std::size_t(a.extent(0)));
  LuFactorizationResult<S> result{.info = {}, .factor = {}};
  std::optional<LuNorms<make_real_t<S>>> norms;
  if (options.retain_norms) norms.emplace();
  auto descriptor = mdspec_of(work);
  dispatch_at_precision<S>(std::forward<Backend>(backend), lu_factor_op{}, p, descriptor, std::span(pivots),
                           norms ? &*norms : nullptr, result.info, pivot_options);
  if (result.info.succeeded())
  {
    uni20::detail::record_output_precision(work, p);
    result.factor = LuFactorAccess::make(std::move(work), std::move(pivots), std::move(norms));
  }
  return result;
}
inline void require_lu_success(SolveInfo const& info)
{
  ERROR_IF(!info.succeeded(), "LU operation failed", static_cast<int>(info.status));
}
} // namespace detail
 
template <KernelBackendSelector Backend, RankedTensorView<2> A>
  requires RealOrComplex<tensor_element_t<A>>
[[nodiscard]] auto lu_factor_with_info(Backend&& backend, A const& a, Precision p,
                                       LuOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  return detail::lu_factor_impl(std::forward<Backend>(backend), a, options, p);
}
 
template <KernelBackendSelector Backend, RankedTensorView<2> A>
  requires RealOrComplex<tensor_element_t<A>>
[[nodiscard]] auto lu_factor_with_info(Backend&& backend, A const& a,
                                       LuOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  return detail::lu_factor_impl(std::forward<Backend>(backend), a, options, precision_of(a));
}
 
template <RankedTensorView<2> A>
  requires RealOrComplex<tensor_element_t<A>>
[[nodiscard]] auto lu_factor_with_info(A const& a, LuOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  auto backend = select_backend_for<DenseMatrix<tensor_element_t<A>>>(lu_factor_op{});
  return lu_factor_with_info(backend, a, options);
}
 
template <RankedTensorView<2> A>
  requires RealOrComplex<tensor_element_t<A>>
[[nodiscard]] auto lu_factor_with_info(A const& a, Precision p,
                                       LuOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  auto backend = select_backend_for<DenseMatrix<tensor_element_t<A>>>(lu_factor_op{});
  return lu_factor_with_info(backend, a, p, options);
}
 
template <class... Args>
  requires requires(Args&&... args) { lu_factor_with_info(std::forward<Args>(args)...); }
[[nodiscard]] auto lu_factor(Args&&... args)
{
  auto result = lu_factor_with_info(std::forward<Args>(args)...);
  detail::require_lu_success(result.info);
  return std::move(*result.factor);
}
 
template <KernelBackendSelector Backend, RealOrComplex S, MutableRankedTensorView<2> B>
  requires std::same_as<S, tensor_element_t<B>>
[[nodiscard]] SolveInfo lu_solve_inplace_with_info(Backend&& backend, LuFactorization<S> const& factor, B&& b)
{
  ERROR_IF(factor.order() != std::size_t(b.extent(0)), "LU solve RHS row count mismatch");
  auto ad = mdspec_of(factor.packed());
  auto bd = mdspec_of(b);
  SolveInfo info;
  auto p = factor.precision();
  detail::dispatch_at_precision<S>(std::forward<Backend>(backend), lu_solve_op{}, p, ad, factor.pivots(), bd, info);
  if (info.succeeded() && factor.order() != 0 && b.extent(1) != 0) uni20::detail::record_output_precision(b, p);
  return info;
}
 
template <RealOrComplex S, MutableRankedTensorView<2> B>
  requires std::same_as<S, tensor_element_t<B>>
[[nodiscard]] SolveInfo lu_solve_inplace_with_info(LuFactorization<S> const& factor, B&& b)
{
  auto backend = select_backend(lu_solve_op{}, factor.packed(), b);
  return lu_solve_inplace_with_info(backend, factor, std::forward<B>(b));
}
 
template <class... Args>
  requires requires(Args&&... args) { lu_solve_inplace_with_info(std::forward<Args>(args)...); }
void lu_solve_inplace(Args&&... args)
{
  detail::require_lu_success(lu_solve_inplace_with_info(std::forward<Args>(args)...));
}
 
template <KernelBackendSelector Backend, RealOrComplex S, RankedTensorView<2> B>
  requires std::same_as<S, tensor_element_t<B>>
[[nodiscard]] auto lu_solve(Backend&& backend, LuFactorization<S> const& factor, B const& b)
{
  auto result = make_tensor<ColumnMajor>(b);
  lu_solve_inplace(std::forward<Backend>(backend), factor, result);
  uni20::detail::record_output_precision(result, factor.precision());
  return result;
}
 
template <RealOrComplex S, RankedTensorView<2> B>
  requires std::same_as<S, tensor_element_t<B>>
[[nodiscard]] auto lu_solve(LuFactorization<S> const& factor, B const& b)
{
  auto backend = select_backend_for<DenseMatrix<S>, DenseMatrix<S>>(lu_solve_op{});
  return lu_solve(backend, factor, b);
}
} // namespace uni20::linalg