QkTable.ipp

#pragma once
#include "CoulombIntegrals.hpp"
#include "IO/ChronoTimer.hpp"
#include "IO/FRW_fileReadWrite.hpp"
#include "Physics/AtomData.hpp"
#include "QkTable.hpp"
#include "fmt/format.hpp"
#include "qip/Widgets.hpp"
#include <algorithm>
#include <cassert>
#include <cstring> // for memcpy
#include <numeric>
#include <string_view>
#include <type_traits>
#include <vector>
 
namespace Coulomb {
 
//==============================================================================
template <Symmetry S>
void CoulombTable<S>::summary() const {
  // const auto symm = to_string(S);
  std::cout << "Stored non-zero two-body integrals (" << to_string(S)
            << " symmetry): \n";
  int k = 0;
  auto total = 0ul;
  for (auto &qk : m_data) {
    fmt::print("k = {:2}: {}\n", k, qk.size());
    // << " [" << qk.bucket_count()<< "]\n";
    total += qk.size();
    ++k;
  }
  std::cout << "Total: " << total << "\n";
}
 
template <Symmetry S>
std::size_t CoulombTable<S>::count() const {
  auto total = 0ul;
  for (auto &qk : m_data) {
    total += qk.size();
  }
  return total;
}
 
//==============================================================================
template <Symmetry S>
void CoulombTable<S>::add(int k, nkIndex a, nkIndex b, nkIndex c, nkIndex d,
                          Real value) {
  add(k, NormalOrder(a, b, c, d), value);
}
//------------------------------------------------------------------------------
template <Symmetry S>
void CoulombTable<S>::add(int k, const DiracSpinor &a, const DiracSpinor &b,
                          const DiracSpinor &c, const DiracSpinor &d,
                          Real value) {
  add(k, a.nk_index(), b.nk_index(), c.nk_index(), d.nk_index(), value);
}
//------------------------------------------------------------------------------
template <Symmetry S>
void CoulombTable<S>::add(int k, nk4Index index, Real value) {
  const auto sk = std::size_t(k);
  if (sk >= m_data.size()) {
    m_data.resize(sk + 1);
  }
  m_data.at(sk).insert({index, value});
}
 
//==============================================================================
template <Symmetry S>
// Updates Q in table. If not present, adds new Q
void CoulombTable<S>::update(int k, nkIndex a, nkIndex b, nkIndex c, nkIndex d,
                             Real value) {
  update(k, NormalOrder(a, b, c, d), value);
}
 
//------------------------------------------------------------------------------
template <Symmetry S>
// Updates Q in table. If not present, adds new Q
void CoulombTable<S>::update(int k, const DiracSpinor &a, const DiracSpinor &b,
                             const DiracSpinor &c, const DiracSpinor &d,
                             Real value) {
  update(k, a.nk_index(), b.nk_index(), c.nk_index(), d.nk_index(), value);
}
 
//------------------------------------------------------------------------------
template <Symmetry S>
void CoulombTable<S>::update(int k, nk4Index index, Real value) {
  const auto sk = std::size_t(k);
  if (sk >= m_data.size()) {
    m_data.resize(sk + 1);
  }
  m_data.at(sk).insert_or_assign(index, value);
}
//==============================================================================
template <Symmetry S>
bool CoulombTable<S>::contains(int k, nkIndex a, nkIndex b, nkIndex c,
                               nkIndex d) const {
  const auto sk = std::size_t(k);
  if (sk >= m_data.size())
    return false;
  return m_data.at(sk).find(NormalOrder(a, b, c, d)) != m_data.at(sk).cend();
}
//------------------------------------------------------------------------------
template <Symmetry S>
bool CoulombTable<S>::contains(int k, const DiracSpinor &a,
                               const DiracSpinor &b, const DiracSpinor &c,
                               const DiracSpinor &d) const {
  return contains(k, a.nk_index(), b.nk_index(), c.nk_index(), d.nk_index());
}
//------------------------------------------------------------------------------
template <Symmetry S>
bool CoulombTable<S>::contains(int k, nk4Index index) const {
  const auto sk = std::size_t(k);
  if (sk >= m_data.size())
    return false;
  return m_data.at(sk).find(index) != m_data.at(sk).cend();
}
 
//==============================================================================
//==============================================================================
 
template <Symmetry S>
double *CoulombTable<S>::get(int k, nk4Index index) {
  const auto sk = std::size_t(k);
  if (sk >= m_data.size())
    return nullptr;
  // check valid k? Probably faster to lookup in table
  const auto map_it = m_data.at(sk).find(index);
  if (map_it == m_data.at(sk).cend())
    return nullptr;
  return &(map_it->second);
}
 
//==============================================================================
template <Symmetry S>
double CoulombTable<S>::Q(int k, nkIndex a, nkIndex b, nkIndex c,
                          nkIndex d) const {
  const auto sk = std::size_t(k);
  if (sk >= m_data.size())
    return 0.0;
  // check valid k? Probably faster to lookup in table
  const auto map_it = m_data.at(sk).find(NormalOrder(a, b, c, d));
  if (map_it == m_data.at(sk).cend())
    return 0.0;
  return map_it->second;
}
//------------------------------------------------------------------------------
template <Symmetry S>
double CoulombTable<S>::Q(int k, const DiracSpinor &a, const DiracSpinor &b,
                          const DiracSpinor &c, const DiracSpinor &d) const {
  return Q(k, a.nk_index(), b.nk_index(), c.nk_index(), d.nk_index());
}
//------------------------------------------------------------------------------
template <Symmetry S>
double CoulombTable<S>::Q(int k, nk4Index index) const {
  const auto sk = std::size_t(k);
  if (sk >= m_data.size())
    return 0.0;
  // check valid k? Probably faster to lookup in table
  const auto map_it = m_data.at(sk).find(index);
  if (map_it == m_data.at(sk).cend())
    return 0.0;
  return map_it->second;
}
 
//==============================================================================
template <Symmetry S>
double CoulombTable<S>::R(int k, nkIndex a, nkIndex b, nkIndex c,
                          nkIndex d) const {
  const auto tQk = Q(k, a, b, c, d);
  if (tQk == 0.0)
    return 0.0;
  const auto s = Angular::neg1pow(k);
 
  const auto [na, ka] = Angular::index_to_nk(a);
  const auto [nb, kb] = Angular::index_to_nk(b);
  const auto [nc, kc] = Angular::index_to_nk(c);
  const auto [nd, kd] = Angular::index_to_nk(d);
 
  const auto tCkac = Angular::tildeCk_kk(k, ka, kc);
  const auto tCkbd = Angular::tildeCk_kk(k, kb, kd);
  return tQk / (s * tCkac * tCkbd);
}
//------------------------------------------------------------------------------
template <Symmetry S>
double CoulombTable<S>::R(int k, const DiracSpinor &a, const DiracSpinor &b,
                          const DiracSpinor &c, const DiracSpinor &d) const {
  return R(k, a.nk_index(), b.nk_index(), c.nk_index(), d.nk_index());
}
 
//==============================================================================
template <Symmetry S>
double CoulombTable<S>::P(int k, nkIndex a, nkIndex b, nkIndex c, nkIndex d,
                          const Angular::SixJTable *const sj) const {
  double Pk_abcd{0.0};
 
  const auto ka = Angular::nkindex_to_kappa(a);
  const auto kb = Angular::nkindex_to_kappa(b);
  const auto kc = Angular::nkindex_to_kappa(c);
  const auto kd = Angular::nkindex_to_kappa(d);
 
  const auto tja = Angular::twoj_k(ka);
  const auto tjb = Angular::twoj_k(kb);
  const auto tjc = Angular::twoj_k(kc);
  const auto tjd = Angular::twoj_k(kd);
 
  // 6j(s) Triads: {a,c,k}, {k,b,d}, {c,b,l}, {d,a,l}
  if (Angular::triangle(tja, tjc, 2 * k) == 0 ||
      Angular::triangle(2 * k, tjb, tjd) == 0)
    return 0.0;
 
  // We may store non-Coulomb integrals in this table (for example, Breit)
  // Then, the parity selection rule doesn't apply,
  // and we don't know what the actual rule is, so use whole range
  const auto [lmin, lmax] = [&]() -> std::pair<int, int> {
    if constexpr (S == Symmetry::Qk)
      return k_minmax_Q(ka, kb, kd, kc);
    else
      return {0, max_k()};
  }();
  const int dl = (S == Symmetry::Qk) ? 2 : 1;
 
  for (int l = lmin; l <= lmax; l += dl) {
    const auto ql = this->Q(l, a, b, d, c); // exchange
    if (ql == 0.0)
      continue;
    const auto sixj = sj ? sj->get_2(tja, tjc, 2 * k, tjb, tjd, 2 * l) :
                           Angular::sixj_2(tja, tjc, 2 * k, tjb, tjd, 2 * l);
    Pk_abcd += sixj * ql;
  }
  Pk_abcd *= double(2 * k + 1);
  return Pk_abcd;
}
 
//------------------------------------------------------------------------------
 
template <Symmetry S>
double CoulombTable<S>::P(int k, const DiracSpinor &a, const DiracSpinor &b,
                          const DiracSpinor &c, const DiracSpinor &d,
                          const Angular::SixJTable *const sj) const {
 
  return P(k, a.nk_index(), b.nk_index(), c.nk_index(), d.nk_index(), sj);
}
 
//==============================================================================
template <Symmetry S>
double CoulombTable<S>::P2(int k, const DiracSpinor &a, const DiracSpinor &b,
                           const DiracSpinor &c, const DiracSpinor &d,
                           const Angular::SixJTable &sj,
                           const std::vector<double> &fk) const {
  double Pk_abcd{0.0};
 
  // 6j(s) Triads: {a,c,k}, {k,b,d}, {c,b,l}, {d,a,l}
  if (Coulomb::triangle(a, c, k) == 0 || Coulomb::triangle(k, b, d) == 0)
    return 0.0;
 
  // We may store non-Coulomb integrals in this table (for example, Breit)
  // Then, the parity selection rule doesn't apply,
  // and we don't know what the actual rule is, so use whole range
  const auto [lmin, lmax] = [&]() -> std::pair<int, int> {
    if constexpr (S == Symmetry::Qk)
      return k_minmax_Q(a, b, d, c);
    else
      return {0, max_k()};
  }();
  const int dl = (S == Symmetry::Qk) ? 2 : 1;
 
  for (int l = lmin; l <= lmax; l += dl) {
    const auto ql = this->Q(l, a, b, d, c); // exchange
    if (ql == 0.0)
      continue;
    const auto sixj = sj.get(a, c, k, b, d, l);
 
    // include effective Coulomb screening:
    const auto f_scr_l = (l < (int)fk.size()) ? fk[std::size_t(l)] : 1.0;
 
    Pk_abcd += f_scr_l * sixj * ql;
  }
  Pk_abcd *= double(2 * k + 1);
  return Pk_abcd;
}
 
//==============================================================================
template <Symmetry S>
double CoulombTable<S>::W(int k, const DiracSpinor &a, const DiracSpinor &b,
                          const DiracSpinor &c, const DiracSpinor &d,
                          const Angular::SixJTable *const sj) const {
  return W(k, a.nk_index(), b.nk_index(), c.nk_index(), d.nk_index(), sj);
}
 
template <Symmetry S>
double CoulombTable<S>::W(int k, nkIndex a, nkIndex b, nkIndex c, nkIndex d,
                          const Angular::SixJTable *const sj) const {
  return Q(k, a, b, c, d) + P(k, a, b, c, d, sj);
}
 
//==============================================================================
template <Symmetry S>
double CoulombTable<S>::g(const DiracSpinor &a, const DiracSpinor &b,
                          const DiracSpinor &c, const DiracSpinor &d, int tma,
                          int tmb, int tmc, int tmd) const {
 
  if (tmc - tma != tmb - tmd)
    return 0.0;
  const int twoq = tmc - tma;
 
  //
  double g = 0.0;
  const auto [k0, ki] = [&]() -> std::pair<int, int> {
    if constexpr (S == Symmetry::Qk)
      return k_minmax_Q(a, b, c, d);
    else
      return {0, max_k()};
  }();
  const int dk = (S == Symmetry::Qk) ? 2 : 1;
 
  for (int k = k0; k <= ki; k += dk) {
    if (std::abs(twoq) > 2 * k)
      continue;
 
    // const auto A =
    //     Angular::neg1pow_2(twoq) *
    //     Angular::Ck_kk_mmq(k, a.kappa(), c.kappa(), tma, tmc, -twoq) *
    //     Angular::Ck_kk_mmq(k, b.kappa(), d.kappa(), tmb, tmd, twoq);
    // if (A != 0.0) {
    //   g += A * R(k, a, b, c, d);
    // }
 
    const auto tjs1 =
      Angular::threej_2(a.twoj(), 2 * k, c.twoj(), -tma, -twoq, tmc);
    const auto tjs2 =
      Angular::threej_2(b.twoj(), 2 * k, d.twoj(), -tmb, twoq, tmd);
    const auto s = Angular::neg1pow_2(2 * k + tmb - tma + twoq);
    g += s * tjs1 * tjs2 * Q(k, a, b, c, d);
  }
  return g;
}
 
//==============================================================================
//==============================================================================
 
//==============================================================================
template <>
inline nk4Index
CoulombTable<Symmetry::Qk>::NormalOrder_impl(nkIndex a, nkIndex b, nkIndex c,
                                             nkIndex d) {
 
  // abcd -> ijkl, with i = min(a,b,c,d)
 
  // abcd = adcb
  // badc = bcda
  // cbad = cdab
  // dabc = dcba
  // nb: there must be a more efficient way of doing this!
 
  // const auto tmp1 = FormIndex(a, b, c, d);
  // const auto tmp2 = FormIndex(a, d, c, b);
  // const auto tmp3 = FormIndex(b, a, d, c);
  // const auto tmp4 = FormIndex(b, c, d, a);
  // const auto tmp5 = FormIndex(c, b, a, d);
  // const auto tmp6 = FormIndex(c, d, a, b);
  // const auto tmp7 = FormIndex(d, a, b, c);
  // const auto tmp8 = FormIndex(d, c, b, a);
  // return std::min({tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7, tmp8});
 
  nk4Index out = FormIndex(a, b, c, d);
  const auto min = std::min({a, b, c, d});
  if (a == min) {
    const auto tmp = FormIndex(a, d, c, b);
    if (tmp < out)
      out = tmp;
  }
  if (b == min) {
    const auto tmp = std::min(FormIndex(b, a, d, c), FormIndex(b, c, d, a));
    if (tmp < out)
      out = tmp;
  }
  if (c == min) {
    const auto tmp = std::min(FormIndex(c, b, a, d), FormIndex(c, d, a, b));
    if (tmp < out)
      out = tmp;
  }
  if (d == min) {
    const auto tmp = std::min(FormIndex(d, a, b, c), FormIndex(d, c, b, a));
    if (tmp < out)
      out = tmp;
  }
  return out;
}
//------------------------------------------------------------------------------
template <>
inline nk4Index
CoulombTable<Symmetry::Wk>::NormalOrder_impl(nkIndex a, nkIndex b, nkIndex c,
                                             nkIndex d) {
  const auto tmp1 = FormIndex(a, b, c, d);
  const auto tmp2 = FormIndex(b, a, d, c);
  const auto tmp3 = FormIndex(c, d, a, b);
  const auto tmp4 = FormIndex(d, c, b, a);
  return std::min({tmp1, tmp2, tmp3, tmp4});
}
 
//------------------------------------------------------------------------------
template <>
inline nk4Index
CoulombTable<Symmetry::Lk>::NormalOrder_impl(nkIndex a, nkIndex b, nkIndex c,
                                             nkIndex d) {
  const auto tmp1 = FormIndex(a, b, c, d);
  const auto tmp2 = FormIndex(b, a, d, c);
  return std::min({tmp1, tmp2});
}
 
//------------------------------------------------------------------------------
template <>
inline nk4Index
CoulombTable<Symmetry::none>::NormalOrder_impl(nkIndex a, nkIndex b, nkIndex c,
                                               nkIndex d) {
  return FormIndex(a, b, c, d);
}
 
//==============================================================================
template <Symmetry S>
nk4Index CoulombTable<S>::NormalOrder(nkIndex a, nkIndex b, nkIndex c,
                                      nkIndex d) const {
  return NormalOrder_impl(a, b, c, d);
}
//------------------------------------------------------------------------------
template <Symmetry S>
nk4Index
CoulombTable<S>::NormalOrder(const DiracSpinor &a, const DiracSpinor &b,
                             const DiracSpinor &c, const DiracSpinor &d) const {
  return NormalOrder(a.nk_index(), b.nk_index(), c.nk_index(), d.nk_index());
}
 
//==============================================================================
template <Symmetry S>
nk4Index CoulombTable<S>::CurrentOrder(nkIndex a, nkIndex b, nkIndex c,
                                       nkIndex d) const {
  return FormIndex(a, b, c, d);
}
//------------------------------------------------------------------------------
template <Symmetry S>
nk4Index CoulombTable<S>::CurrentOrder(const DiracSpinor &a,
                                       const DiracSpinor &b,
                                       const DiracSpinor &c,
                                       const DiracSpinor &d) const {
  return CurrentOrder(a.nk_index(), b.nk_index(), c.nk_index(), d.nk_index());
}
 
//==============================================================================
template <Symmetry S>
nk4Index CoulombTable<S>::FormIndex(nkIndex a, nkIndex b, nkIndex c,
                                    nkIndex d) {
  // create a nk4Index from four small nkIndex, by laying out in memory
 
  // static_assert(sizeof(nk4Index) == 4 * sizeof(nkIndex));
  // nk4Index big_index;
  // const auto p_index = reinterpret_cast<nkIndex *>(&big_index);
  // std::memcpy(p_index, &d, sizeof(nkIndex));
  // std::memcpy(p_index + 1, &c, sizeof(nkIndex));
  // std::memcpy(p_index + 2, &b, sizeof(nkIndex));
  // std::memcpy(p_index + 3, &a, sizeof(nkIndex));
  // return big_index;
 
  // this seems slightly faster...though variance large
  static_assert(sizeof(nk4Index) == 4 * sizeof(nkIndex));
  static_assert(sizeof(nkIndex) * 8 == 16);
  return ((nk4Index)a << 48) + ((nk4Index)b << 32) + ((nk4Index)c << 16) +
         (nk4Index)d;
}
 
//==============================================================================
template <Symmetry S>
std::array<nkIndex, 4>
CoulombTable<S>::UnFormIndex(const nk4Index &index) const {
  static_assert(sizeof(nk4Index) == sizeof(std::array<nkIndex, 4>));
  std::array<nkIndex, 4> set;
  std::memcpy(&set, &index, sizeof(nk4Index));
  // note: this returns {d, c, b, a}, not {a, b, c, d}
  // A warning: this might depend on "endiness"
  // It's OK, this function is never used???
  // If there's an issue, a unit test will fail, so it's "safe"
  std::swap(set[0], set[3]);
  std::swap(set[1], set[2]);
  return set;
}
 
//==============================================================================
template <Symmetry S>
std::array<std::size_t, 4>
CoulombTable<S>::count_non_zero_integrals(const std::vector<DiracSpinor> &basis,
                                          std::size_t k_cutoff,
                                          double eF) const {
  static_assert(S == Symmetry::Qk);
 
  // -1 account for parity selection rule
  // (should be same as 2*max_l, without the -1)
  const auto tmp_max_k =
    std::size_t(std::max(DiracSpinor::max_tj(basis), 1) - 1);
  const auto max_k =
    (k_cutoff <= 0) ? tmp_max_k : std::min(tmp_max_k, std::size_t(k_cutoff));
 
  std::vector<std::size_t> count_non_zero_k_tot(max_k + 1);
  std::vector<std::size_t> count_non_zero_k_Excited(max_k + 1);
  std::vector<std::size_t> count_non_zero_k_MBPT(max_k + 1);
 
  const auto select_Q_Excited = [eF](int, const DiracSpinor &s,
                                     const DiracSpinor &t, const DiracSpinor &u,
                                     const DiracSpinor &v) {
    auto num = Coulomb::number_below_Fermi(s, t, u, v, eF);
    return num == 0;
  };
 
  const auto select_Q_MBPT = [eF](int, const DiracSpinor &s,
                                  const DiracSpinor &t, const DiracSpinor &u,
                                  const DiracSpinor &v) {
    auto num = Coulomb::number_below_Fermi(s, t, u, v, eF);
    return num == 1 || num == 2;
  };
 
#pragma omp parallel for schedule(dynamic, 1)
  for (auto k = 0ul; k <= max_k; ++k) {
    const auto ik = static_cast<int>(k);
    for (const auto &a : basis) {
      for (const auto &b : basis) {
        for (const auto &c : basis) {
          if (!Angular::Ck_kk_SR(ik, a.kappa(), c.kappa()))
            continue;
          for (const auto &d : basis) {
            // due to symmetry, only calculate each 'unique' integral once
            if (NormalOrder(a, b, c, d) == CurrentOrder(a, b, c, d)) {
              if (Angular::Ck_kk_SR(ik, b.kappa(), d.kappa())) {
                ++count_non_zero_k_tot[k];
                if (eF != 0.0) {
                  if (select_Q_Excited(0, a, b, c, d)) {
                    ++count_non_zero_k_Excited[k];
                  }
                  if (select_Q_MBPT(0, a, b, c, d)) {
                    ++count_non_zero_k_MBPT[k];
                  }
                }
              }
            }
          }
        }
      }
    }
  }
  // return count_non_zero_k;
 
  const auto count_non_zero_tot = std::accumulate(
    count_non_zero_k_tot.begin(), count_non_zero_k_tot.end(), 0ul);
 
  const auto count_non_zero_Excited = std::accumulate(
    count_non_zero_k_Excited.begin(), count_non_zero_k_Excited.end(), 0ul);
 
  const auto count_non_zero_MBPT = std::accumulate(
    count_non_zero_k_MBPT.begin(), count_non_zero_k_MBPT.end(), 0ul);
 
  return {count_non_zero_tot, count_non_zero_Excited, count_non_zero_MBPT,
          max_k};
}
 
//==============================================================================
//==============================================================================
template <Symmetry S>
void CoulombTable<S>::fill(const std::vector<DiracSpinor> &basis,
                           const YkTable &yk, int k_cut, bool print) {
  static_assert(S == Symmetry::Qk);
  IO::ChronoTimer t("", print);
 
  /*
  In order to make best use of CPU and Memory, we "fill" the QK table in a
  strange 4-step manner.
  1. Count the number of non-zero Q's for each k
  2. Use info from aboe to 'reserve()' space in the map
  3. Fill all the non-zero Q entries in the map with 0. Note that this adds new
  elements into the map, so cannot be done in parallel. Since we store a
  different map for each k, we may parallelise over k (though this is not super
  efficient)
  4. Now that we have a fully sized map, we can update each of its values in
  parallel
  Note that we make use of the symmetry to ensure we do not access any element
  from more than 1 thread (also, allowing us to not do any more calculations
  than necisary)
  */
 
  // -1 account for parity selection rule
  // (should be same as 2*max_l, without the -1)
  const auto tmp_max_k =
    std::size_t(std::max(DiracSpinor::max_tj(basis), 1) - 1);
 
  const auto max_k =
    (k_cut <= 0) ? tmp_max_k : std::min(tmp_max_k, std::size_t(k_cut));
 
  if (m_data.size() < max_k + 1)
    m_data.resize(max_k + 1);
 
  // 1) Count non-zero Q integrals (each k). Use this to 'reserve' map space=
  std::vector<std::size_t> count_non_zero_k(max_k + 1);
#pragma omp parallel for
  for (auto k = 0ul; k <= max_k; ++k) {
    const auto ik = static_cast<int>(k);
    for (const auto &a : basis) {
      for (const auto &b : basis) {
        for (const auto &c : basis) {
          if (!Angular::Ck_kk_SR(ik, a.kappa(), c.kappa()))
            continue;
          for (const auto &d : basis) {
            // due to symmetry, only calculate each 'unique' integral once
            if (NormalOrder(a, b, c, d) == CurrentOrder(a, b, c, d)) {
              if (Angular::Ck_kk_SR(ik, b.kappa(), d.kappa())) {
                ++count_non_zero_k[k];
              }
            }
          }
        }
      }
    }
  }
 
  // If table already holds every required integral: skip.
  // Pass basis to also check orbital coverage (full fill -> coverage implies
  // completeness; the count alone is fooled by a different basis with a
  // coincidentally-larger per-k count but missing angular-momentum channels).
  if (already_filled(count_non_zero_k, basis)) {
    return;
  }
 
  // 2) Reserve space in each sub-map
#pragma omp parallel for
  for (auto ik = 0ul; ik <= max_k; ++ik) {
    m_data[ik].reserve(count_non_zero_k[ik]);
  }
  const auto n_existing = count();
 
  // 3) Create space in map (set each element to zero).
#pragma omp parallel for schedule(dynamic, 1)
  for (auto k = 0ul; k <= max_k; ++k) {
    for (const auto &a : basis) {
      for (const auto &c : basis) {
        if (!Angular::Ck_kk_SR(int(k), a.kappa(), c.kappa()))
          continue;
        for (const auto &b : basis) {
          for (const auto &d : basis) {
            if (!Angular::Ck_kk_SR(int(k), b.kappa(), d.kappa()))
              continue;
            const auto normal_index = NormalOrder(a, b, c, d);
            if (normal_index == CurrentOrder(a, b, c, d)) {
              m_data[k].insert({normal_index, 0.0});
            }
          }
        }
      }
    }
  }
 
  // 4) Fill the pre-constructed map with values, in parallel. Since we are
  // not adding any new elements to map, and since we are guarenteed to only
  // access each map element once, we can do this part in parallel.
  if (print)
    fmt::print("Calculating {}{} integrals\n",
               n_existing == 0 ? "" : "(subset of) ", count());
  t.start();
  qip::ProgressBar prog(int(basis.size()), print);
#pragma omp parallel for schedule(dynamic, 1)
  for (auto ia = 0ul; ia < basis.size(); ++ia) {
    const auto &a = basis[ia];
    for (const auto &b : basis) {
      for (const auto &c : basis) {
        for (const auto &d : basis) {
 
          const auto normal_index = NormalOrder(a, b, c, d);
          if (normal_index != CurrentOrder(a, b, c, d))
            continue;
 
          const auto [kmin, kmax] = k_minmax_Q(a, b, c, d);
          for (int k = kmin; k <= kmax && k <= int(max_k); k += 2) {
            auto it = m_data[std::size_t(k)].find(normal_index);
            assert(it != m_data[std::size_t(k)].end());
            double *ptr = &it->second;
            if (*ptr == 0.0) {
              // only calculate if not already in table
              // This saves some time, but surprisingly little!
              *ptr = yk.Q(k, a, b, c, d);
            }
          }
        }
      }
    }
    prog.update();
  }
 
  if (print)
    summary();
}
 
//==============================================================================
template <Symmetry S>
void CoulombTable<S>::fill_if(const std::vector<DiracSpinor> &basis,
                              const YkTable &yk,
                              const SelectionRules &SelectionFunction,
                              int k_cut, bool print) {
  static_assert(S == Symmetry::Qk);
  IO::ChronoTimer t("", print);
 
  /*
  Same as above, but only calculate integrals satisfying the SelectionFunction
  */
 
  const auto tmp_max_k =
    std::size_t(std::max(DiracSpinor::max_tj(basis), 1) - 1);
 
  const auto max_k =
    (k_cut <= 0) ? tmp_max_k : std::min(tmp_max_k, std::size_t(k_cut));
 
  if (m_data.size() < max_k + 1)
    m_data.resize(max_k + 1);
 
  // 1) Count non-zero Q integrals (each k). Use this to 'reserve' map space
  std::vector<std::size_t> count_non_zero_k(max_k + 1);
#pragma omp parallel for
  for (auto k = 0ul; k <= max_k; ++k) {
    const auto ik = static_cast<int>(k);
    for (const auto &a : basis) {
      for (const auto &b : basis) {
        for (const auto &c : basis) {
          if (!Angular::Ck_kk_SR(ik, a.kappa(), c.kappa()))
            continue;
          for (const auto &d : basis) {
            // due to symmetry, only calculate each 'unique' integral once
            if (NormalOrder(a, b, c, d) == CurrentOrder(a, b, c, d)) {
              if (Angular::Ck_kk_SR(ik, b.kappa(), d.kappa()) &&
                  SelectionFunction(ik, a, b, c, d)) {
                ++count_non_zero_k[k];
              }
            }
          }
        }
      }
    }
  }
 
  // 2) Reserve space in each sub-map
#pragma omp parallel for
  for (auto ik = 0ul; ik <= max_k; ++ik) {
    m_data[ik].reserve(count_non_zero_k[ik]);
  }
  const auto n_existing = count();
 
  // 3) Create space in map (set each element to zero).
#pragma omp parallel for
  for (auto k = 0ul; k <= max_k; ++k) {
    for (const auto &a : basis) {
      for (const auto &c : basis) {
        if (!Angular::Ck_kk_SR(int(k), a.kappa(), c.kappa()))
          continue;
        for (const auto &b : basis) {
          for (const auto &d : basis) {
            if (!Angular::Ck_kk_SR(int(k), b.kappa(), d.kappa()))
              continue;
            if (!SelectionFunction(int(k), a, b, c, d))
              continue;
            const auto normal_index = NormalOrder(a, b, c, d);
            if (normal_index == CurrentOrder(a, b, c, d)) {
              m_data[k].insert({normal_index, 0.0});
            }
          }
        }
      }
    }
  }
 
  // 4) Fill the pre-constructed map with values, in parallel. Since we are
  // not adding any new elements to map, and since we are guarenteed to only
  // access each map element once, we can do this part in parallel.
  if (print)
    fmt::print("Calculating {}{} integrals\n",
               n_existing == 0 ? "" : "(subset of) ", count());
  t.start();
  qip::ProgressBar prog(int(basis.size()), print);
#pragma omp parallel for schedule(dynamic, 1)
  for (auto ia = 0ul; ia < basis.size(); ++ia) {
    const auto &a = basis[ia];
    for (const auto &b : basis) {
      for (const auto &c : basis) {
        for (const auto &d : basis) {
 
          const auto normal_index = NormalOrder(a, b, c, d);
          if (normal_index != CurrentOrder(a, b, c, d))
            continue;
 
          const auto [kmin, kmax] = k_minmax_Q(a, b, c, d);
          for (int k = kmin; k <= kmax && k <= int(max_k); k += 2) {
            if (!SelectionFunction(k, a, b, c, d))
              continue;
            auto it = m_data[std::size_t(k)].find(normal_index);
            assert(it != m_data[std::size_t(k)].end());
            double *ptr = &it->second;
            if (*ptr == 0.0) {
              *ptr = yk.Q(k, a, b, c, d);
            }
          }
        }
      }
    }
    prog.update();
  }
 
  if (print)
    summary();
}
 
//==============================================================================
template <Symmetry S>
void CoulombTable<S>::fill(const std::vector<DiracSpinor> &basis,
                           const CoulombFunction &Fk,
                           const SelectionRules &Fk_SR, int k_cut, bool print) {
  IO::ChronoTimer t("", print);
 
  /*
  In order to make best use of CPU and Memory, we "fill" the QK table in a
  strange 4-step manner.
  1. Count the number of non-zero Q's for each k
  2. Use info from aboe to 'reserve()' space in the map
  3. Fill all the non-zero Q entries in the map with 0. Note that this adds new
  elements into the map, so cannot be done in parallel. Since we store a
  different map for each k, we may parallelise over k (though this is not super
  efficient)
  4. Now that we have a fully sized map, we can update each of its values in
  parallel
  Note that we make use of the symmetry to ensure we do not access any element
  from more than 1 thread (also, allowing us to not do any more calculations
  than necisary)
  */
 
  const auto tmp_max_k = std::size_t(DiracSpinor::max_tj(basis));
  // May have different parity rule, so don't -1 here
 
  const auto max_k =
    (k_cut <= 0) ? tmp_max_k : std::min(tmp_max_k, std::size_t(k_cut));
 
  if (m_data.size() < max_k + 1)
    m_data.resize(max_k + 1);
 
  // 1) Count non-zero Q integrals (each k). Use this to 'reserve' map space
  std::vector<std::size_t> count_non_zero_k(max_k + 1);
#pragma omp parallel for
  for (auto k = 0ul; k <= max_k; ++k) {
    const auto ik = static_cast<int>(k);
    for (const auto &a : basis) {
      for (const auto &b : basis) {
        for (const auto &c : basis) {
          for (const auto &d : basis) {
            // due to symmetry, only calculate each 'unique' integral once
            if (NormalOrder(a, b, c, d) == CurrentOrder(a, b, c, d)) {
              if (Fk_SR(ik, a, b, c, d)) {
                ++count_non_zero_k[k];
              }
            }
          }
        }
      }
    }
  }
 
  // No early-skip: step 4 only computes entries not already present.
 
  // 2) Reserve space in each sub-map
#pragma omp parallel for
  for (auto ik = 0ul; ik <= max_k; ++ik) {
    m_data[ik].reserve(count_non_zero_k[ik]);
  }
  const auto n_existing = count();
 
  // 3) Create space in map (set each element to zero).
#pragma omp parallel for
  for (auto k = 0ul; k <= max_k; ++k) {
    for (const auto &a : basis) {
      for (const auto &b : basis) {
        for (const auto &c : basis) {
          for (const auto &d : basis) {
            if (Fk_SR(int(k), a, b, c, d)) {
              const auto normal_index = NormalOrder(a, b, c, d);
              if (normal_index == CurrentOrder(a, b, c, d)) {
                m_data[k].insert({normal_index, 0.0});
              }
            }
          }
        }
      }
    }
  }
 
  // 4) Fill the pre-constructed map with values, in parallel. Since we are
  // not adding any new elements to map, and since we are guarenteed to only
  // access each map element once, we can do this part in parallel.
  if (print)
    fmt::print("Calculating {}{} integrals\n",
               n_existing == 0 ? "" : "(subset of) ", count());
  t.start();
 
  qip::ProgressBar prog(int(basis.size()), print);
#pragma omp parallel for schedule(dynamic, 1)
  for (std::size_t ia = 0; ia < basis.size(); ++ia) {
 
    const auto &a = basis[ia];
    for (const auto &b : basis) {
      for (const auto &c : basis) {
        for (const auto &d : basis) {
 
          const auto normal_index = NormalOrder(a, b, c, d);
          if (normal_index != CurrentOrder(a, b, c, d))
            continue;
 
          for (int k = 0; k <= int(max_k); ++k) {
 
            if (!Fk_SR(k, a, b, c, d))
              continue;
            auto it = m_data[std::size_t(k)].find(normal_index);
            assert(it != m_data[std::size_t(k)].end());
            double *ptr = &it->second;
            if (*ptr == 0.0) {
              // only calculate if not already in table
              *ptr = Fk(k, a, b, c, d);
            }
          }
        }
      }
    }
 
    prog.update();
  }
 
  // if (print)
  //   summary();
}
 
//==============================================================================
template <Symmetry S>
void CoulombTable<S>::update(
  const std::vector<DiracSpinor> &basis, const CoulombFunction &Fk, double damp,
  bool print,
  const std::function<bool(const DiracSpinor &, const DiracSpinor &,
                           const DiracSpinor &, const DiracSpinor &)> &filter) {
  IO::ChronoTimer t("", print);
 
  const auto max_k = m_data.size() - 1;
 
  // zero mean no damping (default)
  // damp must be between [0, 1) - cannot equal 1 (meaningless)
  assert(damp != 1.0); // also: catches accidentally passing 'true' (print)
  const auto x = std::clamp(damp, 0.0, 1.0);
  const auto y = 1.0 - damp;
 
  // Fill the pre-constructed map with values, in parallel. Since we are
  // not adding any new elements to map, and since we are guarenteed to only
  // access each map element once, we can do this part in parallel.
 
  if (print) {
    fmt::print("Updating {}{} integrals\n", filter ? "(subset of) " : "",
               count());
  }
  qip::ProgressBar prog(int(basis.size()), print);
#pragma omp parallel for schedule(dynamic, 1)
  for (std::size_t ia = 0; ia < basis.size(); ++ia) {
 
    const auto &a = basis[ia];
    for (const auto &b : basis) {
      for (const auto &c : basis) {
        for (const auto &d : basis) {
 
          // Now, we can check symmetries, selectrion rules etc.
          // just be seeing if this one is in the current table!
          const auto c_index = CurrentOrder(a, b, c, d);
          const auto n_index = NormalOrder(a, b, c, d);
          if (c_index != n_index)
            continue;
 
          // Skip (leave unchanged) entries the filter rejects - before any
          // per-k lookup or recalculation.
          if (filter && !filter(a, b, c, d))
            continue;
 
          for (int k = 0; k <= int(max_k); ++k) {
            double *ptr = get(k, n_index);
            // if not null, means already in table: so update it
            if (ptr != nullptr) {
              *ptr = (*ptr == 0.0 || x == 0.0) ?
                       Fk(k, a, b, c, d) :
                       y * Fk(k, a, b, c, d) + x * *ptr;
            }
          }
        }
      }
    }
 
    prog.update();
  }
}
 
//==============================================================================
// Plain (key, value) pair for bulk (single-call) file read/write of the table
// data. Byte-layout matches the (older) element-by-element format, [key][value]
// pairs in sequence, so files are fully compatible (checked by static_asserts).
struct KeyValueBlock {
  nk4Index key;
  Real value;
};
static_assert(std::is_trivially_copyable_v<KeyValueBlock>);
static_assert(sizeof(KeyValueBlock) == sizeof(nk4Index) + sizeof(Real));
 
//==============================================================================
template <Symmetry S>
void CoulombTable<S>::write(const std::string &fname, bool verbose) const {
  if (fname == "false" || fname == "")
    return;
  if (verbose)
    std::cout << "Writing " << count() << " integrals to file: " << fname
              << ".." << std::flush;
  std::fstream f;
  const auto rw = IO::FRW::write;
  IO::FRW::open_binary(f, fname, rw);
 
  auto size = m_data.size();
  rw_binary(f, rw, size);
  std::vector<KeyValueBlock> block;
  for (const auto &Q_k : m_data) {
    auto size_k = Q_k.size();
    rw_binary(f, rw, size_k);
    // Bulk write: copy block to contiguous buffer, write in single call
    // (element-by-element stream writes dominate the cost otherwise)
    block.clear();
    block.reserve(size_k);
    for (const auto &[key, value] : Q_k) {
      block.push_back({key, value});
    }
    f.write(reinterpret_cast<const char *>(block.data()),
            std::streamsize(block.size() * sizeof(KeyValueBlock)));
  }
  if (verbose)
    std::cout << "\n" << std::flush;
}
 
//==============================================================================
template <Symmetry S>
bool CoulombTable<S>::read(const std::string &fname, bool verbose,
                           bool parse_basis_string) {
  if (fname == "false" || fname.empty())
    return false;
 
  std::fstream f;
  const auto rw = IO::FRW::read;
  IO::FRW::open_binary(f, fname, rw);
 
  // store max_n for each l: allows us to reconstruct the basis string
  // This adds noticable overhead, so only done if parse_basis_string requested
  std::vector<int> max_n_l;
 
  if (!f.good())
    return false;
 
  if (verbose)
    std::cout << "Reading: " << fname << "\n" << std::flush;
 
  std::size_t size{0};
  rw_binary(f, rw, size);
  m_data.resize(size);
  if (m_data.size() == 0)
    return false;
  std::vector<KeyValueBlock> block;
  for (std::size_t i = 0; i < m_data.size(); ++i) {
    auto &Q_k = m_data[i];
    std::size_t size_k{0};
    rw_binary(f, rw, size_k);
    Q_k.reserve(size_k);
 
    // Bulk read: whole k-block in a single call, then insert from memory
    // (element-by-element stream reads dominate the cost otherwise)
    block.resize(size_k);
    f.read(reinterpret_cast<char *>(block.data()),
           std::streamsize(size_k * sizeof(KeyValueBlock)));
    for (const auto &[key, value] : block) {
      Q_k[key] = value;
 
      // determine basis string (max n for each l)
      if (parse_basis_string) {
        const auto indexes = UnFormIndex(key);
        for (const auto index : indexes) {
          const auto [n, ki] = Angular::nkindex_to_n_kindex(index);
          const auto l = std::size_t(Angular::kindex_to_l(ki));
          if (max_n_l.size() < l + 1)
            max_n_l.resize(l + 1);
          if (n > max_n_l.at(l))
            max_n_l.at(l) = n;
        }
      }
    }
  }
 
  if (verbose) {
    // convert list to basis string (max n for each l)
    int last_n = 0;
    int l = 0;
    std::string basis_str;
    for (const auto &n : max_n_l) {
      if (n != last_n)
        basis_str += std::to_string(n);
      basis_str += AtomData::l_symbol(l);
      last_n = n;
      l++;
    }
    if (!basis_str.empty())
      basis_str = "[" + basis_str + "] ";
 
    std::cout << "Read " << count() << " integrals " << basis_str
              << "from file: " << fname << "\n"
              << std::flush;
  }
 
  return true;
}
 
} // namespace Coulomb