tb_hk.hpp

#pragma once
 
#include "./fourier_polynomial.hpp"
#include "./superlattice.hpp"
#include "../../utility/macros.hpp"
 
#include <h5/h5.hpp>
#include <itertools/itertools.hpp>
#include <mpi/mpi.hpp>
#include <nda/h5.hpp>
#include <nda/mpi.hpp>
#include <nda/nda.hpp>
 
#include <algorithm>
#include <array>
#include <ostream>
#include <ranges>
#include <string>
#include <utility>
#include <vector>
 
namespace triqs::experimental::lattice {


  class tb_hk : public fourier_polynomial<2, 3> {
 
    static constexpr auto r_all = nda::range::all;
 
    void check_hoppings();
 
    public:
    using eigenvectors_t = std::pair<nda::array<double, 2>, nda::array<dcomplex, 3>>;

    tb_hk(std::vector<std::array<long, 3>> Rs, std::vector<nda::array<dcomplex, 2>> hoppings)
       : fourier_polynomial<2, 3>(std::move(Rs), std::move(hoppings)) {
 
      check_hoppings();
    };

    tb_hk() : fourier_polynomial<2, 3>({{{0, 0, 0}}}, std::vector(1, nda::array<dcomplex, 2>({{dcomplex(0.)}}))) {};

    C2PY_IGNORE tb_hk(fourier_polynomial<2, 3> fp) : fourier_polynomial<2, 3>{std::move(fp)} {}
 
    // ------------------------ Accessors ----------------------------

    [[nodiscard]] auto hoppings() {
      return std::views::iota(0L, n_R()) | std::views::transform([this](long i) { return coeff_arr(i, nda::ellipsis{}); });
    }

    [[nodiscard]] auto hoppings() const {
      return std::views::iota(0L, n_R()) | std::views::transform([this](long i) { return coeff_arr(i, nda::ellipsis{}); });
    }

    // Intentionally hides the C2PY_IGNORE base accessor to return the list by value, which is what gets wrapped to Python.
    // NOLINTNEXTLINE(bugprone-derived-method-shadowing-base-method)
    [[nodiscard]] std::vector<std::array<long, 3>> get_R_list() const { return fourier_polynomial<2, 3>::get_R_list(); }

    [[nodiscard]] auto elements() { return std::views::zip(get_R_list(), hoppings()); }

    [[nodiscard]] auto elements() const { return std::views::zip(get_R_list(), hoppings()); }

    [[nodiscard]] long n_orbitals() const { return coeff_arr.shape(1); }
 
    //------------------- band basis energy functions ---------------------------

    nda::array<double, 2> eigenvalues(nda::array_view<double, 2> k) const;

    eigenvectors_t eigenvectors(nda::array_view<double, 2> k) const;
 
    // ------------------- Comparison -------------------

    bool operator==(tb_hk const &tb) const {
      return std::ranges::equal(this->get_R_list(), tb.get_R_list()) && this->get_coeff_arr() == tb.get_coeff_arr();
    }
 
    // ------------------- Read / Write -------------------------------

    friend std::ostream &operator<<(std::ostream &out, tb_hk const &tb) {
      out << "tb_hk consisting of " << tb.n_R() << " lattice components (R)";
      out << " with hopping dimension (n_orbitals) " << tb.n_orbitals() << "x" << tb.n_orbitals();
      // if it's a small one, we can print it:
      if (tb.n_R() < 12) {
        out << "\n[";
        for (auto const &[R, C] : tb.elements()) out << "\n   " << R << " : " << C;
        out << " ]";
      }
      return out;
    }

    [[nodiscard]] inline static std::string hdf5_format() { return "tb_hk"; }

    friend void h5_write(h5::group g, std::string const &name, tb_hk const &tb) { tb.h5_write_impl(g, name, "tb_hk"); }

    friend void h5_read(h5::group g, std::string const &name, tb_hk &tb) { tb.h5_read_impl(g, name, "tb_hk"); }
  };

  inline tb_hk fold(superlattice const &sl, tb_hk const &tb) { return fold(sl, static_cast<fourier_polynomial<2, 3> const &>(tb)); }

 
} // namespace triqs::experimental::lattice