gelss_worker.hpp

// 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 "../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 <itertools/itertools.hpp>
 
#include <algorithm>
#include <array>
#include <cmath>
#include <complex>
#include <optional>
#include <utility>
#include <vector>
 
namespace nda::lapack {


  template <typename T>
  class gelss_worker {
    // 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<double, 1> s_;
 
    public:
    int n_var() const { return N_; }

    [[nodiscard]] array<double, 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<double, 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;
      double err = 0.0;
      if (M_ != N_) {
        std::vector<double> 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>, double>{A_plus_ * B, err};
    }

    auto operator()(vector_const_view<T> b, std::optional<long> /*inner_matrix_dim*/ = {}) const {
      using std::sqrt;
      double err = 0.0;
      if (M_ != N_) { err = norm(U_N_H_ * b) / sqrt(b.size()); }
      return std::pair<vector<T>, double>{A_plus_ * b, err};
    }
  };

  struct gelss_worker_hermitian {
    private:
    // Complex double type.
    using dcomplex = std::complex<double>;
 
    // Worker for the original least squares problem.
    gelss_worker<dcomplex> lss_;
 
    // Worker for the extended least squares problem.
    gelss_worker<dcomplex> lss_herm_;
 
    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<dcomplex> A) : lss_(A), lss_herm_(vstack(A, conj(A))) {}

    auto operator()(matrix_const_view<dcomplex> 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<dcomplex, 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<dcomplex>{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<dcomplex>, double>{0.5 * (x + inner_adjoint(x)), err};
    }
  };

 
} // namespace nda::lapack