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 "loaders.hpp"
#include "utils/defs.hpp"
#include "utils/gf_supp.hpp"
#include "utils/nda_supp.hpp"
#include "utils/range_supp.hpp"
#include "utils/embed_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 {
 
    // decomposition of C used by Sigma
    std::vector<long> sigma_embed_decomp;
 
    // decomposition of each impurity
    std::vector<std::vector<long>> imp_decomps;
 
    // names of the sigma (and tau) indices
    std::vector<std::string> _sigma_names;
 
    // names of the alpha indices
    std::vector<std::string> alpha_names;
 
    // stream
    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; 
      long gamma   = 0; 
      long tau     = 0; 
 
      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); }

    [[nodiscard]] nda::array<imp_block_t, 2> psi_map() 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 embed_block_structure() const;

    std::vector<gf_struct_t> imp_block_structure() const;

    std::string description(bool verbosity = false) const;
 
    // ****************** Operation methods to change the embedding **********************

 
    // --------------------------------------------------------
    bool operator==(embedding const &other) const = default;
    // --------------------------------------------------------

    embedding drop_imp(long imp_idx) const;

    embedding replace_imp(long imp_idx_old, long imp_idx_new) const;

    embedding swap_sigma(long alpha) const;

    embedding swap_sigma(std::vector<long> alphas) const;

    embedding slice_sigma() const;

    embedding split_imp(long imp_idx, std::function<bool(long)> p) const;

    embedding split_imp(long imp_idx, std::vector<long> const &block_list) const;
 
    embedding split_imp(long imp_idx, std::initializer_list<const char *> x) = delete;

    embedding split_imp_block(long imp_idx, long gamma, std::vector<long> const &new_dims) const;

    embedding merge_embed_block_by_imp() const;

 
    //------------------------------------------------------------------------------------

    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>>>();
      };
 
      auto make_self_energy = [&](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);
      };
 
      return imp_block_structure() | stdv::transform(make_self_energy) | tl::to<std::vector>();
    }
    //------------------------------------------------------------------------------------

 
    // ****************************************************
    //     Embedding Extract/Embed methods
    // ****************************************************
 
    //--------------------- Extract ---------------------------------------

    template <int Rank> std::vector<std::vector<nda::array<dcomplex, Rank>>> extract(std::vector<nda::array<dcomplex, Rank>> const &X) const {
 
      using array_t = nda::array<dcomplex, Rank>;
 
      // Validate input dimensions
      if (long(X.size()) != n_sigma())
        throw std::runtime_error(
           fmt::format("[embedding::extract] Wrong number of spin channels: got {} but embedding has n_sigma = {}", X.size(), n_sigma()));
 
      auto dim_C = stdr::fold_left(sigma_embed_decomp, 0L, std::plus<>());
      for (long sigma = 0; sigma < n_sigma(); ++sigma) {
        auto actual_dim = X[sigma].extent(detail::has_frequency<Rank> ? 1 : 0);
        if (actual_dim != dim_C)
          throw std::runtime_error(fmt::format(
             "[embedding::extract] Dimension mismatch for spin channel {}: got {} but expected C-space dimension {}", sigma, actual_dim, dim_C));
      }
 
      // Does this rank have a frequency index?
      [[maybe_unused]] long n_w = 0;
      if constexpr (detail::has_frequency<Rank>) { n_w = X[0].extent(0); }
 
      // Extract from the array X the blocks for each (alpha, sigma)
      auto embed_blocks = nda::array<array_t, 2>(n_alpha(), n_sigma());
      for (auto &&[alpha, r_alpha] : enumerated_sub_slices(sigma_embed_decomp)) {
        for (auto sigma : range(n_sigma())) { embed_blocks(alpha, sigma) = detail::extract_diagonal(X[sigma], r_alpha); }
      }
 
      // Lambda to extract one impurity using the reverse_psi map
      auto imp_gf_stru_list = imp_block_structure();
      auto extract_one_imp  = [&](long n_imp) -> std::vector<array_t> {
        std::vector<array_t> blocks_imp;
        for (auto [bl_name, bl_size] : imp_gf_stru_list[n_imp]) { blocks_imp.push_back(detail::make_zero_array<Rank, dcomplex>(bl_size, n_w)); }
        auto const &rpsi = reverse_psi[n_imp];
        for (auto [gamma, tau] : rpsi.indices()) {
          auto [alpha, sigma]                      = rpsi(gamma, tau)[0];
          blocks_imp[gamma + n_gamma(n_imp) * tau] = embed_blocks(alpha, sigma);
        }
        return blocks_imp;
      };
 
      // Apply lambda to each impurity.
      return range(n_impurities()) | stdv::transform(extract_one_imp) | tl::to<std::vector>();
    }
 
    //--------------------- Embed -----------------------------------------

    template <int Rank> std::vector<nda::array<dcomplex, Rank>> embed(std::vector<std::vector<nda::array<dcomplex, Rank>>> const &imps_blocks) const {
 
      using array_t = nda::array<dcomplex, Rank>;
 
      // Validate input structure
      if (long(imps_blocks.size()) != n_impurities())
        throw std::runtime_error(
           fmt::format("[embedding::embed] Wrong number of impurities: got {} but embedding has {}", imps_blocks.size(), n_impurities()));
 
      auto imp_gf_stru = imp_block_structure();
      for (long imp = 0; imp < n_impurities(); ++imp) {
        auto expected_n_blocks = long(imp_gf_stru[imp].size());
        if (long(imps_blocks[imp].size()) != expected_n_blocks)
          throw std::runtime_error(
             fmt::format("[embedding::embed] Impurity {} has {} blocks but embedding expects {}", imp, imps_blocks[imp].size(), expected_n_blocks));
        for (long b = 0; b < expected_n_blocks; ++b) {
          auto expected_dim = imp_gf_stru[imp][b].second;
          auto actual_dim   = imps_blocks[imp][b].extent(detail::has_frequency<Rank> ? 1 : 0);
          if (actual_dim != expected_dim)
            throw std::runtime_error(
               fmt::format("[embedding::embed] Impurity {}, block {} has dimension {} but embedding expects {}", imp, b, actual_dim, expected_dim));
        }
      }
 
      // Does this rank have a frequency index?
      [[maybe_unused]] long n_w = 0;
      if constexpr (detail::has_frequency<Rank>) { n_w = imps_blocks[0][0].extent(0); }
 
      // Build the embedded array
      auto embed_blocks = nda::array<array_t, 2>(n_alpha(), n_sigma());
      for (auto &&[alpha, sigma] : psi.indices()) {
        auto bl_size               = sigma_embed_decomp[alpha];
        embed_blocks(alpha, sigma) = detail::make_zero_array<Rank, dcomplex>(bl_size, n_w);
      }
 
      // Use psi map to fill the embedded blocks array
      for (auto &&[block, m] : zip(embed_blocks, psi)) {
        if (m.imp_idx == -1) continue;
        block = imps_blocks[m.imp_idx][m.gamma + n_gamma(m.imp_idx) * m.tau];
      }
 
      // Total dimension of the C space
      auto dim_C = stdr::fold_left(sigma_embed_decomp, 0, std::plus<>());
 
      // Build the result from the embedded view and return
      std::vector<array_t> result;
      for (auto sigma : range(n_sigma())) {
        auto full_tensor = detail::make_zero_array<Rank, dcomplex>(dim_C, n_w);
        for (auto &&[alpha, r_alpha] : enumerated_sub_slices(sigma_embed_decomp)) {
          detail::embed_diagonal(full_tensor, r_alpha, embed_blocks(alpha, sigma));
        }
        result.push_back(std::move(full_tensor));
      }
 
      return result;
    }

    template <typename Mesh> block2_gf<Mesh, matrix_valued> embed(std::vector<block_gf<Mesh, matrix_valued>> const &Sigma_imp_vec) const {
      // Check number of impurities
      if (long(Sigma_imp_vec.size()) != n_impurities())
        throw std::runtime_error(fmt::format("[embedding::embed] Wrong number of impurity self-energies: got {} but embedding has {} impurities",
                                             Sigma_imp_vec.size(), n_impurities()));
      // 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"};
 
      // build the result
      auto const &mesh = Sigma_imp_vec[0][0].mesh();
      auto Sigma_embed = make_block2_gf(mesh, this->embed_block_structure());
      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::vector<block_gf<Mesh, matrix_valued>> extract(block2_gf<Mesh, matrix_valued> const &g_loc) const {
      // Check that the decomposition of g_loc matches either sigma_embed_decomp
      if (auto decomp = get_struct(g_loc).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_loc, this->embed_block_structure()));
      }
 
      auto imp_gf_stru_list = imp_block_structure();
      auto extract_one_imp  = [&](long n_imp) {
        auto gimp        = block_gf{g_loc(0, 0).mesh(), imp_gf_stru_list[n_imp]};
        auto const &rpsi = reverse_psi[n_imp];
        for (auto [gamma, tau] : rpsi.indices()) {
          if (rpsi(gamma, tau).empty()) continue;
          auto [alpha, sigma]                       = rpsi(gamma, tau)[0];
          gimp[gamma + n_gamma(n_imp) * tau].data() = g_loc(alpha, sigma).data();
        }
        return gimp;
      };
 
      return range(n_impurities()) | stdv::transform(extract_one_imp) | tl::to<std::vector>();
    }

    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 {
      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 {embed(Sigma_imp_vec), embed(Sigma_imp_static_vec)};
    }

    block2_matrix_t embed(std::vector<block_matrix_t> const &Sigma_imp_static_vec) const;

    std::vector<block_matrix_t> extract(block2_matrix_t const &matrix_C) const;
  };
 
  // ***************************** Free functions/factories *******************************************


  embedding make_embedding(std::vector<std::string> const &spin_names, nda::array<std::vector<long>, 2> const &block_decomposition,
                           std::vector<long> const &atom_to_imp);

  embedding make_embedding(std::vector<std::string> const &spin_names, std::vector<std::vector<long>> const &block_decomposition,
                           std::vector<long> const &atom_to_imp);

  embedding make_embedding(local_space const &C_space, bool use_atom_equivalences = true, bool use_atom_decomp = false);
 
  // -----------------------------------------------------------------------
  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};
  }

 
  template std::vector<std::vector<nda::array<dcomplex, 2>>> embedding::extract(std::vector<nda::array<dcomplex, 2>> const &) const;
  template std::vector<std::vector<nda::array<dcomplex, 3>>> embedding::extract(std::vector<nda::array<dcomplex, 3>> const &) const;
  template std::vector<std::vector<nda::array<dcomplex, 4>>> embedding::extract(std::vector<nda::array<dcomplex, 4>> const &) const;
  template std::vector<std::vector<nda::array<dcomplex, 5>>> embedding::extract(std::vector<nda::array<dcomplex, 5>> const &) const;
 
  template std::vector<nda::array<dcomplex, 2>> embedding::embed(std::vector<std::vector<nda::array<dcomplex, 2>>> const &) const;
  template std::vector<nda::array<dcomplex, 3>> embedding::embed(std::vector<std::vector<nda::array<dcomplex, 3>>> const &) const;
  template std::vector<nda::array<dcomplex, 4>> embedding::embed(std::vector<std::vector<nda::array<dcomplex, 4>>> const &) const;
  template std::vector<nda::array<dcomplex, 5>> embedding::embed(std::vector<std::vector<nda::array<dcomplex, 5>>> const &) const;
 
#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