embedding.hpp

Go to the documentation of this file.

// Copyright (c) 2025--present, The Simons Foundation
// This file is part of TRIQS/modest and is licensed under the terms of GPLv3 or later.
// SPDX-License-Identifier: GPL-3.0-or-later
// See LICENSE in the root of this distribution for details.
 
#pragma once
#include "./downfolding.hpp"
#include "utils/defs.hpp"
#include "utils/gf_supp.hpp"
#include "utils/range_supp.hpp"
#include <fmt/ranges.h>
 
namespace triqs {
  using block_matrix_t  = std::vector<nda::matrix<dcomplex>>;
  using block2_matrix_t = nda::array<nda::matrix<dcomplex>, 2>;
} // namespace triqs
 
namespace triqs::modest {
 
  class embedding {
 
    std::vector<long> sigma_embed_decomp;
 
    std::vector<std::vector<long>> imp_decomps;
 
    std::vector<std::string> _sigma_names;
 
    std::vector<std::string> alpha_names;
 
    friend std::ostream &operator<<(std::ostream &out, embedding const &E);
 
    friend void h5_write(h5::group g, std::string const &name, embedding const &x);
    friend void h5_read(h5::group g, std::string const &name, embedding &x);
 
    public:
    struct imp_block_t {
 
      long imp_idx = 0; // index of the impurity
      long gamma   = 0; // first index γ of the block in the Σ_imp
      long tau     = 0; // second index τ of the block in the Σ_imp
 
      imp_block_t() = default;
      imp_block_t(long n_imp, long gamma, long tau) : imp_idx(n_imp), gamma(gamma), tau(tau) {}
      friend bool operator==(imp_block_t const &, imp_block_t const &) = default;
 
      // FIXME : merge with format_as
      friend std::ostream &operator<<(std::ostream &out, imp_block_t const &x) {
        return out << fmt::format("(imp_idx = {}, γ = {}, τ = {})", x.imp_idx, x.gamma, x.tau);
      }
    };
    private:
    // For each imp, for all (γ, τ), a list of all (α, σ) such that ψ[α, σ] ==(n_imp, γ, τ)
    nda::array<imp_block_t, 2> psi; // ψ [α, σ] --> (n_imp, γ, τ)
 
    // reverse ψ [n_imp] -> (γ, τ) -> [(α, σ)]
    std::vector<nda::array<std::vector<std::array<long, 2>>, 2>> reverse_psi;
 
    public:
 
    embedding(std::vector<long> sigma_embed_decomposition, std::vector<std::vector<long>> imp_decompositions, nda::array<imp_block_t, 2> psi,
              std::vector<std::string> sigma_names);
 
    embedding() = default; // for h5_read
 
 
    imp_block_t operator[](long alpha, long sigma) const { return psi(alpha, sigma); }
 
    // HL: get_psi -> psi to match other classes?
    [[nodiscard]] nda::array<imp_block_t, 2> get_psi() const { return psi; }
 
    [[nodiscard]] long n_alpha() const { return psi.extent(0); }
 
    [[nodiscard]] long n_sigma() const { return psi.extent(1); }
 
    [[nodiscard]] long n_gamma(long imp_idx) const { return long(imp_decomps[imp_idx].size()); }
 
    [[nodiscard]] long n_impurities() const { return long(imp_decomps.size()); }
 
    [[nodiscard]] std::vector<std::string> sigma_names() const { return _sigma_names; };
 
    [[nodiscard]] std::vector<long> imp_decomposition(long imp_idx) const { return imp_decomps[imp_idx]; };
 
    C2PY_IGNORE gf_struct2_t sigma_embed_block_shape() const;
 
    std::vector<gf_struct_t> imp_block_shape() const;
 
    std::string description(bool verbosity = false) const;
 
    // ****************** Operation methods to change the embedding **********************
 
 
    // --------------------------------------------------------
    bool operator==(embedding const &other) const = default;
    // --------------------------------------------------------
 
    embedding drop(long imp_idx) const;
 
    embedding replace(long imp_idx_to_remove, long imp_idx_to_replace_with) const;
 
    embedding flip_spin(long alpha) const;
 
    embedding flip_spin(std::vector<long> alphas) const;
 
    embedding split(long imp_idx, std::function<bool(long)> p) const;
 
    embedding split(long imp_idx, std::vector<long> const &block_list) const;
 
    embedding split(long imp_idx, std::initializer_list<const char *> x) = delete;
 
    //------------------------------------------------------------------------------------
    template <typename Mesh> std::vector<std::pair<block_gf<Mesh, matrix_valued>, block_matrix_t>> make_zero_imp_self_energies(Mesh const &mesh) {
      auto make_block_matrix = [](auto const &gf_struct) {
        return gf_struct | stdv::transform([](auto &x) {
                 auto bl_size = x.second;
                 return nda::zeros<dcomplex>(bl_size, bl_size);
               })
           | tl::to<std::vector<nda::matrix<dcomplex>>>();
      };
      return imp_block_shape() | stdv::transform([&](auto const &gf_struct) {
               auto Sigma_static  = make_block_matrix(gf_struct);
               auto Sigma_dynamic = block_gf<Mesh, matrix_valued>{mesh, gf_struct};
               return std::make_pair(Sigma_dynamic, Sigma_static);
             })
         | tl::to<std::vector>();
    }
    //------------------------------------------------------------------------------------
 
    // ****************************************************
    //     Embedding Extract/Embed methods
    // ****************************************************
    //--------------------- Embed -----------------------------------------
 
 
    template <typename Mesh> block2_gf<Mesh, matrix_valued> embed(std::vector<block_gf<Mesh, matrix_valued>> const &Sigma_imp_vec) const;
 
    template <typename Mesh>
    std::pair<block2_gf<Mesh, matrix_valued>, block2_matrix_t> embed(std::vector<block_gf<Mesh, matrix_valued>> const &Sigma_imp_vec,
                                                                     std::vector<block_matrix_t> const &Sigma_imp_static_vec) const;
 
    block2_matrix_t embed(std::vector<block_matrix_t> const &Sigma_imp_static_vec) const;
 
    nda::array<nda::array<dcomplex, 3>, 2> embed(std::vector<std::vector<nda::array<dcomplex, 3>>> const &Sigma_imp_vec) const;
 
    std::vector<nda::array<dcomplex, 5>> embed(std::vector<nda::array<dcomplex, 5>> const &pi_imp_vec) const;
 
    // nda::array<dcomplex, 4> embed(std::vector<nda::array<dcomplex, 4>> const &U_tensor_vec) const;
 
    //--------------------- Extract ---------------------------------------
 
    template <typename Mesh> std::vector<block_gf<Mesh, matrix_valued>> extract(block2_gf<Mesh, matrix_valued> const &g_loc) const;
 
    std::vector<block_matrix_t> extract(block2_matrix_t const &matrix_C) const;
 
    std::vector<std::vector<nda::array<dcomplex, 3>>> extract(nda::array<dcomplex, 4> const &g_loc) const;
 
    std::vector<nda::array<dcomplex, 5>> extract(nda::array<dcomplex, 5> const &Pi_loc) const;
 
    // std::vector<nda::array<dcomplex, 4>> extract(nda::array<dcomplex, 4> const &U_tensor) const;
  };
 
  // ***************************** Free functions/factories *******************************************
 
  embedding make_embedding_impl(local_space const &C_space, nda::array<std::vector<long>, 2> const &block_decomposition,
                                std::optional<std::vector<long>> const &atom_to_imp = std::nullopt);
  embedding make_embedding(local_space const &C_space, bool use_atom_equivalences = true);
 
  // *****************************  Method implementation ***************************************
 
  template <typename Mesh> block2_gf<Mesh, matrix_valued> embedding::embed(std::vector<block_gf<Mesh, matrix_valued>> const &Sigma_imp_vec) const {
    // Check that all meshes are the same for Sigma_imp_vec
    if (not all_equal(Sigma_imp_vec | stdv::transform([](auto &&x) -> decltype(auto) { return x[0].mesh(); })))
      throw std::runtime_error{"[embedding_desc::embed]: meshes of solvers are not all equal"};
 
    auto const &mesh = Sigma_imp_vec[0][0].mesh();
 
    // build the result
    auto Sigma_embed = make_block2_gf(mesh, this->sigma_embed_block_shape());
    for (auto &&[S, m] : zip(Sigma_embed, psi)) {
      if (m.imp_idx == -1) continue;
      S() = Sigma_imp_vec[m.imp_idx][m.gamma + n_gamma(m.imp_idx) * m.tau];
    }
    return Sigma_embed;
  }
 
  template <typename Mesh>
  std::pair<block2_gf<Mesh, matrix_valued>, block2_matrix_t> embedding::embed(std::vector<block_gf<Mesh, matrix_valued>> const &Sigma_imp_vec,
                                                                              std::vector<block_matrix_t> const &Sigma_imp_static_vec) const {
    if (Sigma_imp_vec.size() != Sigma_imp_static_vec.size()) {
      throw std::runtime_error(fmt::format("The lists of self-energies are not equal {} != {}", Sigma_imp_vec.size(), Sigma_imp_static_vec.size()));
    }
    return {this->embed(Sigma_imp_vec), this->embed(Sigma_imp_static_vec)};
  }
 
  // ---------------------------------------------------------------
 
  template <typename Mesh> std::vector<block_gf<Mesh, matrix_valued>> embedding::extract(block2_gf<Mesh, matrix_valued> const &g) const {
 
    if (auto decomp = get_struct(g).dims(r_all, 0) | tl::to<std::vector>(); decomp != this->sigma_embed_decomp) {
      if (decomp.size() != 1) throw std::runtime_error{"extract: g should have decomp = sigma_embedding_decomp or [1]"};
      return extract(decomposition_view(g, this->sigma_embed_block_shape()));
    }
 
    // FIXME : check all meshes are the same
    auto imp_gf_stru_list = imp_block_shape();
    auto extract_one_imp  = [&](long n_imp) {
      auto gimp        = block_gf{g(0, 0).mesh(), imp_gf_stru_list[n_imp]};
      auto const &rpsi = reverse_psi[n_imp];
      for (auto [gamma, tau] : rpsi.indices()) {
        auto [alpha, sigma]                       = rpsi(gamma, tau)[0];
        gimp[gamma + n_gamma(n_imp) * tau].data() = g(alpha, sigma).data();
      }
      return gimp;
    };
    return range(n_impurities()) | stdv::transform(extract_one_imp) | tl::to<std::vector>();
  }
 
  // -----------------------------------------------------------------------
  inline std::pair<one_body_elements_on_grid, embedding> make_embedding_with_clusters(one_body_elements_on_grid obe,
                                                                                      std::vector<std::vector<long>> const &atom_partition) {
    auto new_obe = permute_local_space(atom_partition, obe);
    auto E       = make_embedding(new_obe.C_space, false);
    return {new_obe, E};
  }
 
  // ------------------------------------------------------------------------------
 
#define INSTANTIATE(Mesh)                                                                                                                            \
  template block2_gf<Mesh, matrix_valued> embedding::embed(std::vector<block_gf<Mesh, matrix_valued>> const &) const;                                \
  template std::pair<block2_gf<Mesh, matrix_valued>, block2_matrix_t> embedding::embed(std::vector<block_gf<Mesh, matrix_valued>> const &,           \
                                                                                       std::vector<block_matrix_t> const &) const;                   \
  template std::vector<block_gf<Mesh, matrix_valued>> embedding::extract(block2_gf<Mesh, matrix_valued> const &) const;                              \
  template std::vector<std::pair<block_gf<Mesh, matrix_valued>, block_matrix_t>> embedding::make_zero_imp_self_energies(Mesh const &);
  INSTANTIATE(imfreq);
  INSTANTIATE(refreq);
  INSTANTIATE(dlr_imfreq);
#undef INSTANTIATE
 
} // namespace triqs::modest