ops/least_squares.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/linalg/backends/cpu/least_squares.hpp>
#include <uni20/linalg/backends/lapack/least_squares.hpp>
#include <uni20/tensor/precision.hpp>
 
namespace uni20::linalg
{
template <RealOrComplex S> struct LeastSquaresResult
{
    LeastSquaresInfo info = {};
    std::optional<DenseMatrix<S>> solution = {};
    std::vector<make_real_t<S>> residual_norms = {};
};
 
template <KernelBackendSelector Backend, RankedTensorView<2> A, RankedTensorView<2> B>
  requires RealOrComplex<tensor_element_t<A>> && std::same_as<tensor_element_t<A>, tensor_element_t<B>>
[[nodiscard]] auto least_squares_with_info(Backend&& backend, A const& a, B const& b, Precision p,
                                           LeastSquaresOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  using S = tensor_element_t<A>;
  require_precision<S>(p);
  (void)p.bit_count();
  ERROR_IF(a.extent(0) < a.extent(1) || a.extent(0) != b.extent(0),
           "least squares requires m >= n and matching RHS rows");
  least_squares_detail::require_options(options);
  LeastSquaresResult<S> result;
  result.solution.emplace(uninitialized, a.extent(1), b.extent(1), p);
  result.residual_norms.resize(b.extent(1));
  auto ad = mdspec_of(a);
  auto bd = mdspec_of(b);
  auto xd = mdspec_of(*result.solution);
  detail::dispatch_at_precision<S>(std::forward<Backend>(backend), least_squares_op{}, p, xd, ad, bd,
                                   std::span(result.residual_norms), result.info, options);
  if (!result.info.succeeded())
  {
    result.solution.reset();
    result.residual_norms.clear();
  }
  return result;
}
 
template <KernelBackendSelector Backend, RankedTensorView<2> A, RankedTensorView<2> B>
  requires RealOrComplex<tensor_element_t<A>> && std::same_as<tensor_element_t<A>, tensor_element_t<B>>
[[nodiscard]] auto least_squares_with_info(Backend&& backend, A const& a, B const& b,
                                           LeastSquaresOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  return least_squares_with_info(std::forward<Backend>(backend), a, b, common_default_precision(a, b), options);
}
 
template <RankedTensorView<2> A, RankedTensorView<2> B>
  requires RealOrComplex<tensor_element_t<A>> && std::same_as<tensor_element_t<A>, tensor_element_t<B>>
[[nodiscard]] auto least_squares_with_info(A const& a, B const& b, Precision p,
                                           LeastSquaresOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  return least_squares_with_info(select_backend_for<DenseMatrix<tensor_element_t<A>>>(least_squares_op{}), a, b, p,
                                 options);
}
 
template <RankedTensorView<2> A, RankedTensorView<2> B>
  requires RealOrComplex<tensor_element_t<A>> && std::same_as<tensor_element_t<A>, tensor_element_t<B>>
[[nodiscard]] auto least_squares_with_info(A const& a, B const& b,
                                           LeastSquaresOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  return least_squares_with_info(a, b, common_default_precision(a, b), options);
}
 
template <class... Args>
  requires requires(Args&&... args) { least_squares_with_info(std::forward<Args>(args)...); }
[[nodiscard]] auto least_squares(Args&&... args)
{
  auto result = least_squares_with_info(std::forward<Args>(args)...);
  ERROR_IF(!result.info.succeeded(), "least squares failed", static_cast<int>(result.info.status));
  return std::move(*result.solution);
}
} // namespace uni20::linalg