local_space.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 <triqs/gfs.hpp>
#include <triqs/utility/macros.hpp>
#include <mpi/mpi.hpp>
#include "./utils/defs.hpp"
#include "./utils/gf_supp.hpp"
 
namespace triqs::modest {
 
  using namespace triqs::gfs;
 
  // -----------------------------------------------
  enum class spin_kind_e {
    Polarized,    // σ = 0,1, the object can have different value for different σ
    NonPolarized, // σ = 0,1, but the object is the same for both σ. Can store only one copy e.g.
    NonColinear   // σ = 0. There is no index, the spin index is grouped with other non-diagonal indices. e.g. Nambu, spin-orbit
  };
 
  inline long n_sigma_from_spin_kind(spin_kind_e sk) {
    switch (sk) {
      case spin_kind_e::Polarized: return 2;
      case spin_kind_e::NonPolarized: return 2;
      case spin_kind_e::NonColinear: return 1;
    }
    __builtin_unreachable();
  }
 
  // --------------------------------------------------------------------
  struct atomic_orbs {
    long dim     = 0; 
    long l       = 0; 
    long cls_idx = 0; 
    long dft_idx = 0; 
    // int n_irrep;     // are the orbitals reducible into irreps (irep) ( seach "irep" in TRIQS/dft_tools; never used in DMFT routines
    void serialize(auto &ar) const { ar & dim & l & cls_idx & dft_idx; };
    void deserialize(auto &ar) { ar & dim & l & cls_idx & dft_idx; };
 
    bool operator==(atomic_orbs const &) const = default;
 
    C2PY_IGNORE friend void mpi_broadcast(atomic_orbs &x, mpi::communicator c = {}, int root = 0) {
      mpi::broadcast(x.dim, c, root);
      mpi::broadcast(x.l, c, root);
      mpi::broadcast(x.cls_idx, c, root);
      mpi::broadcast(x.dft_idx, c, root);
    }
  };
 
  // ==========================================================
  class local_space {
 
    spin_kind_e _spin_kind = spin_kind_e::NonPolarized;
 
    // List of all atomic shells spanning the 𝓒 space
    std::vector<atomic_orbs> _atomic_shells;
 
    // For each atom, a list of the dimensions of the irreps decomposition of the atomic shell
    nda::array<std::vector<long>, 2> _irreps_decomp_per_atom;
 
    // rotation matrices that rotate BACK to the basis of the DFT code basis.
    // FIXME : CHECK vs INVERSE, EXPLAIN, WRITE NOTE.
    nda::array<nda::matrix<dcomplex>, 2> _rotation_from_dft_to_local_basis;
 
    // rotation matrix from the Ylm basis to the dft code
    nda::array<nda::matrix<dcomplex>, 1> _rotation_from_spherical_to_dft_basis;
 
    // dimension of the correlated space
    long _dim_C = 0;
 
    // atom names
    std::vector<std::string> _atom_names = {};
 
    friend void h5_read(h5::group g, std::string const &name, local_space &ls);
    friend void h5_write(h5::group g, std::string const &name, local_space const &ls);
 
    C2PY_IGNORE friend void mpi_broadcast(local_space &x, mpi::communicator c = {}, int root = 0) {
      mpi::broadcast(x._spin_kind, c, root);
      mpi::broadcast(x._atomic_shells, c, root);
      mpi::broadcast(x._irreps_decomp_per_atom, c, root);
      mpi::broadcast(x._rotation_from_dft_to_local_basis, c, root);
      mpi::broadcast(x._rotation_from_spherical_to_dft_basis, c, root);
      mpi::broadcast(x._dim_C, c, root);
      mpi::broadcast(x._atom_names, c, root);
    }
 
    public:
    bool operator==(local_space const &) const = default;
 
    local_space(spin_kind_e spin_kind, std::vector<atomic_orbs> atomic_shells, nda::array<std::vector<long>, 2> irreps_decomp_per_atom,
                nda::array<nda::matrix<dcomplex>, 2> rotation_from_dft_to_local_basis,
                nda::array<nda::matrix<dcomplex>, 1> rotation_from_spherical_to_dft_basis);
 
    local_space() = default;
    // ---------------- sigma indices and co. ----------------
 
 
    [[nodiscard]] spin_kind_e spin_kind() const { return _spin_kind; };
 
    [[nodiscard]] long n_sigma() const { return n_sigma_from_spin_kind(_spin_kind); }
 
    [[nodiscard]] std::vector<std::string> sigma_names() const {
      return (_spin_kind == spin_kind_e::NonColinear ? std::vector{"ud"s} : std::vector{"up"s, "down"s});
    }
 
    // --------  Atomic shells and related functions -----------
    [[nodiscard]] long dim() const { return _dim_C; }
 
    [[nodiscard]] std::vector<atomic_orbs> const &atomic_shells() const { return _atomic_shells; }
 
    [[nodiscard]] long n_atoms() const { return long(_atomic_shells.size()); }
 
    // OP : explain the 2 indices : spin ?? [ a, σ] ?
    [[nodiscard]] nda::array<std::vector<long>, 2> const &atoms_block_decomposition() const { return _irreps_decomp_per_atom; }
 
    [[nodiscard]] nda::array<nda::matrix<dcomplex>, 2> const &rotation_from_dft_to_local_basis() const { return _rotation_from_dft_to_local_basis; }
 
    [[nodiscard]] nda::array<nda::matrix<dcomplex>, 1> const &rotation_from_spherical_to_dft_basis() const {
      return _rotation_from_spherical_to_dft_basis;
    }
 
    [[nodiscard]] long first_shell_of_its_equiv_cls(long idx) const;
 
    [[nodiscard]] std::vector<std::string> atom_names() const { return _atom_names; }
 
    [[nodiscard]] auto atomic_decomposition() const {
      return atomic_shells() | stdv::transform([](auto const &s) { return s.dim; });
    }
 
 
    // --------  Atomic "view" methods for Green functions and matrices -----------
 
    // FIXME : IT IS NOT A VIEW, IT IS A COPY
    // Where do we use this ? do we want a view or a copy ?
 
    template <typename Mesh> block2_gf<Mesh> atomic_view(block2_gf<Mesh> const &G_C) {
      // FIXME : ASSERT check the dimension of G_C m x m , one a block, n_sigma ...
      auto Gout    = make_block2_gf(G_C(0, 0).mesh(), this->Gatom_block_shape());
      auto n_sigma = G_C.size2();
      for (auto sigma : range(n_sigma))
        for (auto const &[atom, r_atom] : enumerated_sub_slices(this->atomic_decomposition())) {
          Gout(atom, sigma).data() = G_C(0, sigma).data()(r_all, r_atom, r_atom);
        }
      return Gout;
    }
 
    nda::array<nda::matrix<dcomplex>, 2> atomic_view(nda::array<nda::matrix<dcomplex>, 2> const &matrix_C);
 
    // --------  gf_structs for different block shapes -----------
 
    [[nodiscard]] C2PY_IGNORE gf_struct2_t Gc_block_shape() const;
 
    // OP : [a, \sigma] : what is a, sigma ?
    [[nodiscard]] C2PY_IGNORE gf_struct2_t Gatom_block_shape() const {
      auto res = nda::zeros<long>(_atomic_shells.size(), n_sigma());
      for (auto const &[alpha, shell] : enumerate(_atomic_shells)) res(alpha, r_all) = shell.dim;
      return {.names = {atom_names(), sigma_names()}, .dims = std::move(res)};
    }
 
  };
 
  // --------  stream -----------
 
  std::ostream &operator<<(std::ostream &out, local_space const &bd);
 
  // --------  instantiations --------------
 
  template block2_gf<imfreq> local_space::atomic_view(block2_gf<imfreq> const &G_C);
} // namespace triqs::modest