tail_fitter.hpp

Go to the documentation of this file.

// Copyright (c) 2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
// Copyright (c) 2018 Centre national de la recherche scientifique (CNRS)
// Copyright (c) 2018-2023 Simons Foundation
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You may obtain a copy of the License at
//     https://www.gnu.org/licenses/gpl-3.0.txt
//
// Authors: Olivier Parcollet, Nils Wentzell

 
#pragma once
 
#include "../arrays.hpp"
#include "../utility/macros.hpp"
 
#include <itertools/itertools.hpp>
#include <nda/nda.hpp>
#include <nda/lapack/gelss_worker.hpp>
 
#include <algorithm>
#include <array>
#include <cmath>
#include <complex>
#include <memory>
#include <optional>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace triqs::mesh {
  // Forward declaration.
  class imfreq;
} // namespace triqs::mesh
 
namespace triqs::mesh::detail {


  inline auto vander(std::vector<std::complex<double>> const &z_pts, int q) {
    nda::matrix<std::complex<double>> V(z_pts.size(), q + 1);
    for (auto [i, z_i] : itertools::enumerate(z_pts)) {
      auto z = std::complex<double>{1};
      for (int n = 0; n <= q; ++n) {
        V(i, n) = z;
        z *= z_i;
      }
    }
    return V;
  }

  template <int R> auto tail_eval(nda::array_const_view<std::complex<double>, R> A, std::complex<double> z_0) {
    auto compute = [&A, z_0](auto res) {
      auto z       = std::complex<double>{1};
      auto const q = A.extent(0);
      for (int n = 0; n < q; ++n, z /= z_0) res += A(n, nda::ellipsis{}) * z;
      return res;
    };
    if constexpr (R > 1) {
      // return an array of rank R - 1
      return compute(nda::zeros<std::complex<double>>(nda::stdutil::front_pop(A.shape())));
    } else {
      // return a complex scalar
      return compute(std::complex<double>{0});
    }
  }
 
} // namespace triqs::mesh::detail
 
namespace triqs::mesh {

  class C2PY_IGNORE tail_fitter {
    public:
    static constexpr double default_tail_fraction = 0.2;

    static constexpr int default_n_tail_max = 30;

    tail_fitter(double r, int p_max, std::optional<int> q = {}) : r_(r), p_max_(p_max), adjust_q_(not q.has_value()), q_(adjust_q_ ? q_max_ : *q) {
      if (r_ <= 0 or r_ > 1) TRIQS_RUNTIME_ERROR << "Error in tail-fitter: Fraction of mesh points must be in (0, 1]";
      if (p_max_ <= 0) TRIQS_RUNTIME_ERROR << "Error in tail-fitter: Maximum number of mesh points must be > 0";
    }

    template <typename M> int n_pts_in_tail(M const &m) const { return std::min(static_cast<int>(std::round(r_ * m.size() / 2)), p_max_); }

    double get_tail_fraction() const { return r_; }

    template <typename M> auto get_tail_fit_indices(M const &m) {
      // total number of points in the fitting window
      auto const p_r = static_cast<int>(std::round(r_ * m.size() / 2));
 
      // number of points actually used for the fit
      auto const p = n_pts_in_tail(m);
 
      // reserve space for the indices
      std::vector<long> idx_vec;
      idx_vec.reserve(2ul * p);
 
      // initialize the left most and right most indices for both fitting windows
      double const step = static_cast<double>(p_r) / p;
      double left_idx   = m.first_index();
      double right_idx  = m.last_index();
 
      for ([[maybe_unused]] auto i : nda::range(p)) {
        idx_vec.push_back(long(left_idx));
        idx_vec.push_back(long(right_idx));
        left_idx += step;
        right_idx -= step;
      }
 
      return idx_vec;
    }

    template <bool enforce_hermiticity = false> auto &get_lss() {
      if constexpr (enforce_hermiticity)
        return lss_hermitian_;
      else
        return lss_;
    }

    template <bool enforce_hermiticity = false, typename M> void setup_lss(M const &m, int n_A) {
      // least square worker type
      using worker_t = std::conditional_t<enforce_hermiticity, nda::lapack::gelss_worker_hermitian, nda::lapack::gelss_worker<std::complex<double>>>;
 
      // indices of the mesh points to use in the tail fit
      if (fit_idxs_.empty()) fit_idxs_ = get_tail_fit_indices(m);
 
      // set up Vandermonde matrix (the points are given by z_i = |m.w_max()| / m.to_value(n))
      double const z_max = std::abs(m.w_max());
      if (V_.is_empty()) {
        std::vector<std::complex<double>> z_pts;
        z_pts.reserve(fit_idxs_.size());
        for (long n : fit_idxs_) z_pts.push_back(z_max / m.to_value(n));
        V_ = detail::vander(z_pts, q_);
      }
 
      // check if we have enough data points for the least square procedure (p > n_A + 1)
      if (n_A + 1 > V_.extent(0) / 2) TRIQS_RUNTIME_ERROR << "Error in tail-fitter::setup_lss: Insufficient data points for least square procedure";
 
      // factory function for least square workers
      auto worker_factory = [&](int n) { return std::make_unique<const worker_t>(V_(nda::range::all, nda::range(n_A, n + 1))); };
 
      // get the correct (hermitian vs. non-hermitian) array of least square workers
      auto &lss = get_lss<enforce_hermiticity>();
 
      // set up the least square workers
      if (!adjust_q_) {
        // use the expansion order given in the constructor
        lss[n_A] = worker_factory(q_);
      } else {
        // find the maximum expansion order such that the smallest singular value of the Vandermonde matrix is > rcond_
        lss[n_A].reset();
        // ensure that |z_max|^{1-q} > 10^{-16}
        long q_max = std::min(static_cast<long>(q_max_), static_cast<long>(1. + 16. / std::log10(1 + std::abs(m.w_max()))));
        // we try to use at least two times as many data points as we have unknown coefficients
        q_max = std::min(q_max, V_.extent(0) / 2);
        for (long q = q_max; q >= n_A; --q) {
          auto ptr = worker_factory(q);
          if (ptr->S_vec()[ptr->S_vec().size() - 1] > rcond_) {
            lss[n_A] = std::move(ptr);
            break;
          }
        }
      }
 
      // throw an exception if the Vandermonde matrix is ill-conditioned
      if (!lss[n_A]) TRIQS_RUNTIME_ERROR << "Error in tail-fitter::setup_lss: Ill-conditioned Vandermonde matrix";
    }

    template <int P, bool enforce_hermiticity = false, typename M, int R>
    auto fit(M const &m, nda::array_const_view<std::complex<double>, R> D, bool rescale, nda::array_const_view<std::complex<double>, R> C,
             std::optional<long> d = {}) {
      // compile-time and run-time checks
      static_assert(!enforce_hermiticity || std::is_same_v<M, imfreq>);
      if (enforce_hermiticity and not d.has_value())
        TRIQS_RUNTIME_ERROR << "Error in tail-fitter::fit: Enforcing hermiticity requires an inner matrix dimension";
      if (m.positive_only()) TRIQS_RUNTIME_ERROR << "Error in tail-fitter::fit: Cannot fit on a positive_only mesh";
 
      // early return if the number of known coefficients is larger than the expansion order
      int const n_A = C.extent(0);
      if (n_A > q_) return std::pair<nda::array<std::complex<double>, R>, double>{C, 0.0};
 
      // set up the least squares worker for the given number of known coefficients if it has not been done already
      auto &lss = get_lss<enforce_hermiticity>();
      if (!lss[n_A]) setup_lss<enforce_hermiticity>(m, n_A);
 
      // permute the indices of D such that the relevant frequency mesh corresponds to the first dimension
      auto D_rot = nda::rotate_index_view<P>(D);
 
      // flatten D in the target space and the remaining meshes into the second dimension of a new matrix
      long const ncols = D_rot.size() / D_rot.shape()[0];
      auto D_mat       = nda::matrix<std::complex<double>>(V_.extent(0), ncols);
      for (auto [i, n] : itertools::enumerate(fit_idxs_)) {
        if constexpr (R == 1) {
          D_mat(i, 0) = D_rot(m.to_data_index(n));
        } else {
          for (auto [j, x] : itertools::enumerate(D_rot(m.to_data_index(n), nda::ellipsis{}))) { D_mat(i, j) = x; }
        }
      }
 
      // flatten and prepare the array of known coefficients C
      double const z_max = std::abs(m.w_max());
      if (n_A > 0) {
        // check the shape of C
        if (ncols != C.size() / C.shape()[0])
          TRIQS_RUNTIME_ERROR << "Error in tail-fitter::fit: Shape of C array incompatible with the shape of the D array";
 
        // flatten C and scale its values by |z_max|^{-q}
        double z   = 1.0;
        auto C_mat = nda::matrix<std::complex<double>>(n_A, ncols);
        for (long i : nda::range(n_A)) {
          if constexpr (R == 1) {
            C_mat(i, 0) = z * C(i, nda::ellipsis{});
          } else {
            for (auto [j, x] : itertools::enumerate(C(i, nda::ellipsis{}))) C_mat(i, j) = z * x;
          }
          z /= z_max;
        }
 
        // subtract the expansion terms corresponding to the known moments from the function values
        D_mat -= V_(nda::range::all, nda::range(n_A)) * C_mat;
      }
 
      // perform the least squares procedure
      auto [A_mat, err] = (*lss[n_A])(D_mat, d);
 
      // rescale the coefficients if requested
      if (rescale) {
        double z = 1.0;
        for ([[maybe_unused]] long i : nda::range(n_A)) z *= z_max;
        for (long i : nda::range(A_mat.extent(0))) {
          A_mat(i, nda::range::all) *= z;
          z *= z_max;
        }
      }
 
      // reinterpret the result as an R-dimensional array according to the initial shape
      auto shape = D_rot.shape();
      shape[0]   = lss[n_A]->n_var() + n_A;
      auto A     = nda::array<std::complex<double>, R>{shape};
 
      // add the known moments to the result
      if (n_A) A(nda::range(n_A), nda::ellipsis{}) = C;
 
      // add the calculated moments to the result
      shape[0]    = lss[n_A]->n_var();
      auto idxmap = typename nda::array<std::complex<double>, R>::layout_t{shape};
      auto A_view = A(nda::range(n_A, A.shape()[0]), nda::ellipsis{});
      A_view      = nda::array_view<std::complex<double>, R>{idxmap, A_mat.storage()};
 
      return std::pair<nda::array<std::complex<double>, R>, double>{std::move(A), err};
    }

    template <int P, typename M, int R>
    auto fit_hermitian(M const &m, nda::array_const_view<std::complex<double>, R> D, bool rescale, nda::array_const_view<std::complex<double>, R> C,
                       std::optional<long> d = {}) {
      return fit<P, true, M, R>(m, D, rescale, C, d);
    }
 
    private:
    static constexpr int q_max_    = 9;
    static constexpr double rcond_ = 1e-8;
    double r_;
    int p_max_;
    bool adjust_q_;
    int q_;
    std::array<std::unique_ptr<const nda::lapack::gelss_worker<dcomplex>>, q_max_ + 1> lss_;
    std::array<std::unique_ptr<const nda::lapack::gelss_worker_hermitian>, q_max_ + 1> lss_hermitian_;
    nda::matrix<dcomplex> V_;
    std::vector<long> fit_idxs_;
  };

  class C2PY_IGNORE tail_fitter_handle {
    public:
    void set_tail_fit_parameters(double tail_fraction, int n_tail_max = tail_fitter::default_n_tail_max,
                                 std::optional<int> expansion_order = {}) const {
      tf_ptr_ = std::make_shared<tail_fitter>(tail_fitter{tail_fraction, n_tail_max, expansion_order});
    }

    C2PY_IGNORE tail_fitter &get_tail_fitter() const {
      if (!tf_ptr_) tf_ptr_ = std::make_shared<tail_fitter>(tail_fitter::default_tail_fraction, tail_fitter::default_n_tail_max);
      return *tf_ptr_;
    }

    C2PY_IGNORE tail_fitter &get_tail_fitter(double tail_fraction, int n_tail_max = tail_fitter::default_n_tail_max,
                                             std::optional<int> expansion_order = {}) const {
      set_tail_fit_parameters(tail_fraction, n_tail_max, expansion_order);
      return *tf_ptr_;
    }
 
    private:
    mutable std::shared_ptr<tail_fitter> tf_ptr_;
  };

 
} // namespace triqs::mesh