gelss_worker.hpp

Go to the documentation of this file.

// Copyright (c) 2020--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.

 
#pragma once
 
#include "./gesvd.hpp"
#include "../algorithms.hpp"
#include "../arithmetic.hpp"
#include "../basic_array.hpp"
#include "../declarations.hpp"
#include "../exceptions.hpp"
#include "../layout/policies.hpp"
#include "../layout_transforms.hpp"
#include "../linalg.hpp"
#include "../mapped_functions.hpp"
#include "../matrix_functions.hpp"
#include "../traits.hpp"
 
#include <itertools/itertools.hpp>
 
#include <algorithm>
#include <array>
#include <cmath>
#include <complex>
#include <optional>
#include <utility>
#include <vector>
 
namespace nda::lapack {


  template <typename T>
    requires is_blas_lapack_v<T>
  class gelss_worker {
    public:
    using value_type = T;

    using fp_type = get_fp_t<T>;

    [[nodiscard]] int n_var() const { return N_; }

    [[nodiscard]] array<fp_type, 1> const &S_vec() const { return s_; }

    gelss_worker(matrix_const_view<T> A) : M_(A.extent(0)), N_(A.extent(1)), s_(std::min(M_, N_)) {
      if (N_ > M_) NDA_RUNTIME_ERROR << "Error in nda::lapack::gelss_worker: Matrix A cannot have more columns than rows";
 
      // initialize matrices
      matrix<T, F_layout> A_work{A};
      matrix<T, F_layout> U(M_, M_);
      matrix<T, F_layout> V_H(N_, N_);
 
      // calculate the SVD: A = U * \Sigma * V^H
      gesvd(A_work, s_, U, V_H);
 
      // calculate the pseudo inverse A^{+} = V * \Sigma^{+} * U^H
      matrix<fp_type, F_layout> S_plus(N_, M_);
      S_plus = 0.;
      for (long i : range(s_.size())) S_plus(i, i) = 1.0 / s_(i);
      A_plus_ = dagger(V_H) * S_plus * dagger(U);
 
      // set U_N^H
      if (N_ < M_) U_N_H_ = dagger(U)(range(N_, M_), range(M_));
    }

    auto operator()(matrix_const_view<T> B, std::optional<long> /* inner_matrix_dim */ = {}) const {
      using std::sqrt;
      fp_type err = 0.0;
      if (M_ != N_) {
        std::vector<fp_type> err_vec;
        for (long i : range(B.shape()[1])) err_vec.push_back(frobenius_norm(U_N_H_ * B(range::all, range(i, i + 1))) / sqrt(B.shape()[0]));
        err = *std::ranges::max_element(err_vec);
      }
      return std::pair<matrix<T>, fp_type>{A_plus_ * B, err};
    }

    auto operator()(vector_const_view<T> b, std::optional<long> /*inner_matrix_dim*/ = {}) const {
      using std::sqrt;
      fp_type err = 0.0;
      if (M_ != N_) err = nda::linalg::norm(U_N_H_ * b) / sqrt(b.size());
      return std::pair<vector<T>, fp_type>{A_plus_ * b, err};
    }
 
    private:
    // Number of rows (M) and columns (N) of the Matrix A.
    long M_, N_;
 
    // Pseudo inverse of A, i.e. A^{+} = V * \Sigma^{+} * U^H.
    matrix<T> A_plus_;
 
    // U_N^H defining the error of the least squares problem.
    matrix<T> U_N_H_;
 
    // Array containing the singular values.
    array<fp_type, 1> s_;
  };

  class gelss_worker_hermitian {
    public:
    int n_var() const { return lss_herm_.n_var(); }

    [[nodiscard]] array<double, 1> const &S_vec() const { return lss_.S_vec(); }

    gelss_worker_hermitian(matrix_const_view<std::complex<double>> A) : lss_(A), lss_herm_(vstack(A, conj(A))) {}

    auto operator()(matrix_const_view<std::complex<double>> B, std::optional<long> inner_matrix_dim = {}) const {
      if (not inner_matrix_dim.has_value())
        NDA_RUNTIME_ERROR << "Error in nda::lapack::gelss_worker_hermitian: Inner matrix dimension required for hermitian least square fitting";
      long d = *inner_matrix_dim;
 
      // take the inner 'adjoint' of a matrix C:
      // * reshape C -> C': (M, N) -> (M, N', d, d)
      // * for each m and n: C'(m, n, :, :) -> C'(m, n, :, :)^/dagger
      // * reshape C' -> C: (M, N', d, d) -> (M, N)
      auto inner_adjoint = [d](auto &C) {
        NDA_ASSERT2(C.shape()[1] % (d * d) == 0, "Error in nda::lapack::gelss_worker_hermitian: Data shape incompatible with inner matrix dimension");
        auto shape = C.shape();
 
        // get extent N' of second dimension
        long N = shape[1] / (d * d);
 
        // reshape, transpose and take the complex conjugate
        array<std::complex<double>, 4> arr_dag =
           conj(permuted_indices_view<encode(std::array{0, 1, 3, 2})>(reshape(C, std::array{shape[0], N, d, d})));
 
        // return the result in a new matrix
        return matrix<std::complex<double>>{reshape(std::move(arr_dag), shape)};
      };
 
      // solve the extended system vstack(A, A*) * X = vstack(B, B_dag)
      auto B_dag    = inner_adjoint(B);
      auto [x, err] = lss_herm_(vstack(B, B_dag));
 
      // resymmetrize the results to cure small hermiticity violations
      return std::pair<matrix<std::complex<double>>, double>{0.5 * (x + inner_adjoint(x)), err};
    }
 
    private:
    // Worker for the original least squares problem.
    gelss_worker<std::complex<double>> lss_;
 
    // Worker for the extended least squares problem.
    gelss_worker<std::complex<double>> lss_herm_;
  };

 
} // namespace nda::lapack