brzone.hpp

Go to the documentation of this file.

// Copyright (c) 2016-2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
// Copyright (c) 2016-2018 Centre national de la recherche scientifique (CNRS)
// Copyright (c) 2018-2023 Simons Foundation
// Copyright (c) 2023 Hugo U.R. Strand
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You may obtain a copy of the License at
//     https://www.gnu.org/licenses/gpl-3.0.txt
//
// Authors: Thomas Ayral, Philipp Dumitrescu, Dominik Kiese, Olivier Parcollet, Hugo U. R. Strand, Nils Wentzell

 
#pragma once
 
#include "./k_expr.hpp"
#include "./mesh_iterator.hpp"
#include "./utils.hpp"
#include "../lattice/brillouin_zone.hpp"
#include "../utility/macros.hpp"
 
#include <fmt/ranges.h>
#include <h5/h5.hpp>
#include <nda/nda.hpp>
 
#include <array>
#include <cmath>
#include <cstdint>
#include <concepts>
#include <iostream>
#include <limits>
#include <mutex>
#include <optional>
#include <ranges>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace triqs::mesh {

  using lattice::brillouin_zone;


  class C2PY_RENAME(MeshBrZone) brzone {
    public:
    using value_t = brillouin_zone::value_t;

    using index_t = std::array<long, 3>;

    using data_index_t = long;

    class C2PY_IGNORE mesh_point_t {
      public:
      using mesh_t = brzone;

      mesh_point_t() = default;

      mesh_point_t(std::array<long, 3> const &n, brzone const *m_ptr, long d)
         : index_(n), m_ptr_(m_ptr), data_index_(d), mesh_hash_(m_ptr->mesh_hash()) {}

      mesh_point_t(mesh_point_t const &mp)
         : index_(mp.index_), m_ptr_(mp.m_ptr_), data_index_(mp.data_index_), mesh_hash_(mp.mesh_hash_), value_(mp.value_) {}

      [[nodiscard]] mesh_t::index_t index() const { return index_; }

      [[nodiscard]] long data_index() const { return data_index_; }

      [[nodiscard]] value_t const &value() const {
        if (value_)
          return *value_;
        else {
          auto guard = std::lock_guard{value_mutex_};
          return *(value_ = m_ptr_->to_value(index_));
        }
      }

      [[nodiscard]] uint64_t mesh_hash() const noexcept { return mesh_hash_; }

      [[nodiscard]] operator value_t() const { return value(); }

      [[deprecated("() is deprecated for a brzone::mesh_point_t. Use [] instead")]] double operator()(int d) const { return value()[d]; }

      [[nodiscard]] double operator[](int i) const { return value()[i]; }

      friend std::ostream &operator<<(std::ostream &sout, mesh_point_t const &mp) { return sout << mp.value(); }
 
      private:
      std::array<long, 3> index_            = {0, 0, 0};
      brzone const *m_ptr_                  = nullptr;
      long data_index_                      = 0;
      uint64_t mesh_hash_                   = 0;
      mutable std::optional<value_t> value_ = {};
      mutable std::mutex value_mutex_       = {};
    };
 
    public:
    brzone() = default;

    brzone(brillouin_zone const &bz, std::array<long, 3> const &dims)
       : bz_(bz),
         dims_(dims),
         size_(nda::stdutil::product(dims)),
         s2_(dims_[2]),
         s1_(dims_[1] * dims_[2]),
         units_(nda::linalg::inv(1.0 * nda::diag(dims)) * bz.units()),
         units_inv_(nda::linalg::inv(units_)),
         mesh_hash_(hash(nda::sum(bz.units()), dims[0], dims[1], dims[2])) {
      EXPECTS(dims_[0] > 0 and dims_[1] > 0 and dims_[2] > 0);
    }

    C2PY_IGNORE brzone(brillouin_zone const &bz, nda::matrix<long> const &M) : brzone(bz, std::array{M(0, 0), M(1, 1), M(2, 2)}) {
      EXPECTS((M.shape() == std::array{3l, 3l}));
      EXPECTS(nda::is_matrix_diagonal(M));
    }

    brzone(brillouin_zone const &bz, long n_k) : brzone(bz, std::array{n_k, (bz.ndim() >= 2 ? n_k : 1l), (bz.ndim() >= 3 ? n_k : 1)}) {}

    [[nodiscard]] bool is_index_valid(index_t const &n) const noexcept {
      for (auto i : nda::range(3))
        if (n[i] < 0 or n[i] >= dims_[i]) return false;
      return true;
    }

    [[nodiscard]] data_index_t to_data_index(index_t const &n) const {
      EXPECTS(is_index_valid(n));
      return n[0] * s1_ + n[1] * s2_ + n[2];
    }

    template <typename V> [[nodiscard]] data_index_t to_data_index(closest_mesh_point_t<V> const &cmp) const {
      return to_data_index(closest_index(cmp.value));
    }

    template <char OP, typename L> [[nodiscard]] data_index_t to_data_index(k_expr_unary<OP, L> const &ex) const {
      EXPECTS(mesh_hash_ == ex.mesh_hash());
      return to_data_index(index_modulo(ex.index()));
    }

    template <char OP, typename L, typename R> [[nodiscard]] data_index_t to_data_index(k_expr<OP, L, R> const &ex) const {
      EXPECTS(mesh_hash_ == ex.mesh_hash());
      return to_data_index(index_modulo(ex.index()));
    }

    [[nodiscard]] index_t to_index(data_index_t d) const {
      EXPECTS(0 <= d and d < size());
      long const r0 = d % s1_;
      return {d / s1_, r0 / s2_, r0 % s2_};
    }

    template <typename V> [[nodiscard]] index_t to_index(closest_mesh_point_t<V> const &cmp) const { return closest_index(cmp.value); }

    [[nodiscard]] mesh_point_t operator[](long d) const { return {to_index(d), this, d}; }

    [[nodiscard]] C2PY_IGNORE mesh_point_t operator[](closest_mesh_point_t<value_t> const &cmp) const { return (*this)[this->to_data_index(cmp)]; }

    [[nodiscard]] mesh_point_t operator()(index_t const &n) const { return {n, this, to_data_index(n)}; }

    [[nodiscard]] value_t to_value(index_t const &n) const {
      EXPECTS(is_index_valid(n));
      return nda::transpose(units_)(nda::range::all, nda::range(bz_.ndim())) * nda::basic_array_view{n}(nda::range(bz_.ndim()));
    }

    [[nodiscard]] C2PY_PROPERTY_GET(dims) auto const &dims() const { return dims_; }

    [[nodiscard]] C2PY_PROPERTY_GET(units) auto units() const { return nda::matrix_const_view<double>{units_}; }

    [[nodiscard]] C2PY_PROPERTY_GET(units_inv) auto units_inv() const { return nda::matrix_const_view<double>{units_inv_}; }

    [[nodiscard]] C2PY_PROPERTY_GET(bz) auto const &bz() const noexcept { return bz_; }

    [[nodiscard]] C2PY_PROPERTY_GET(mesh_hash) uint64_t mesh_hash() const { return mesh_hash_; }

    [[nodiscard]] long size() const { return size_; }

    [[nodiscard]] index_t index_modulo(index_t const &n_tilde) const {
      return {positive_modulo(n_tilde[0], dims_[0]), positive_modulo(n_tilde[1], dims_[1]), positive_modulo(n_tilde[2], dims_[2])};
    }

    template <typename V>
      requires(is_k_expr<V> or std::ranges::contiguous_range<V> or nda::ArrayOfRank<V, 1>)
    [[nodiscard]] index_t closest_index(V const &k) const {
      if constexpr (is_k_expr<V>) {
        return closest_index(k.value());
      } else {
        // calculate k in the brzone basis
        auto ks      = nda::stack_vector<double, 3>{k[0], k[1], k[2]};
        auto k_units = nda::transpose(units_inv_) * ks;
        auto n       = nda::stack_vector<long, 3>(nda::floor(k_units));
 
        // calculate position relative to neighbors in mesh
        auto w = k_units - n;
 
        // prepare result container and distance measure
        auto dst = std::numeric_limits<double>::infinity();
 
        // check flatness along mesh dimensions
        long r1 = std::min(dims_[0], 2l);
        long r2 = std::min(dims_[1], 2l);
        long r3 = std::min(dims_[2], 2l);
 
        // find nearest neighbor by comparing distances
        nda::stack_vector<long, 3> res;
        for (auto const &[i1, i2, i3] : itertools::product_range(r1, r2, r3)) {
          auto iv   = nda::stack_vector<long, 3>{i1, i2, i3};
          auto dstp = nda::linalg::norm(nda::transpose(units_) * (w - iv));
 
          // update result when distance is smaller than current
          if (dstp < dst) {
            dst = dstp;
            res = n + iv;
          }
        }
 
        // fold back to brzone mesh (nearest neighbor could be out of bounds)
        return index_modulo({res[0], res[1], res[2]});
      }
    }

    [[nodiscard]] auto begin() const { return mesh_iterator<brzone>{.mesh_ptr = this, .data_index = 0}; }

    [[nodiscard]] auto cbegin() const { return begin(); }

    [[nodiscard]] auto end() const { return mesh_iterator<brzone>{.mesh_ptr = this, .data_index = size()}; }

    [[nodiscard]] auto cend() const { return end(); }

    bool operator==(brzone const &m) const { return bz_ == m.bz() && dims_ == m.dims(); }

    friend std::ostream &operator<<(std::ostream &sout, brzone const &m) {
      return sout << "Brillouin zone mesh with " << m.dims() << " k-points and an underlying " << m.bz();
    }

    void serialize(auto &ar) const { ar & bz_ & dims_ & size_ & s2_ & s1_ & units_ & units_inv_ & mesh_hash_; }

    void deserialize(auto &ar) { ar & bz_ & dims_ & size_ & s2_ & s1_ & units_ & units_inv_ & mesh_hash_; }

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

    friend void h5_write(h5::group g, std::string const &name, brzone const &m) {
      h5::group gr = g.create_group(name);
      h5::write_hdf5_format(gr, m); // NOLINT (downcasting to base class)
      h5::write(gr, "dims", m.dims_);
      h5::write(gr, "brillouin_zone", m.bz_);
    }

    friend void h5_read(h5::group g, std::string const &name, brzone &m) {
      h5::group gr = g.open_group(name);
      h5::assert_hdf5_format(gr, m, true); // NOLINT (downcasting to base class)
 
      std::array<long, 3> dims{};
      if (gr.has_key("dims")) {
        h5::read(gr, "dims", dims);
      } else {
        // for backward compatibility
        auto M = h5::read<matrix<long>>(gr, "periodization_matrix");
        dims   = {M(0, 0), M(1, 1), M(2, 2)};
      }
 
      brillouin_zone bz{};
      if (gr.has_key("brillouin_zone")) {
        h5::read(gr, "brillouin_zone", bz);
      } else if (gr.has_key("bz")) {
        // for backward compatibility
        h5::read(gr, "bz", bz);
      } else {
        std::cout << "WARNING: Reading old MeshBrillouinZone without BrillouinZone\n";
      }
 
      m = brzone(bz, dims);
    }
 
    private:
    brillouin_zone bz_             = {};
    std::array<long, 3> dims_      = {0, 0, 0};
    long size_                     = 0;
    long s2_                       = 1;
    long s1_                       = 1;
    nda::matrix<double> units_     = nda::eye<double>(3);
    nda::matrix<double> units_inv_ = nda::eye<double>(3);
    uint64_t mesh_hash_            = 0;
  };
 
  namespace detail {
 
    // Helper struct to evaluate a function at an arbitrary k-vector.
    struct brzone1d {
      long dim;
    };
 
  } // namespace detail
 
  // Helper function to do linear interpolation between k-points.
  auto evaluate(detail::brzone1d const &m, auto const &f, double vi) {
    long i   = static_cast<long>(std::floor(vi));
    double w = vi - double(i);
    return (1 - w) * f(positive_modulo(i, m.dim)) + w * f(m.dim == 1 ? 0 : positive_modulo(i + 1, m.dim));
  }

  auto evaluate(brzone const &m, auto const &f, brzone::index_t const &n_tilde) { return f(m.index_modulo(n_tilde)); }

  template <typename V>
    requires(std::ranges::contiguous_range<V> or nda::ArrayOfRank<V, 1> or is_k_expr<V>)
  auto evaluate(brzone const &m, auto const &f, V const &k) {
    if constexpr (is_k_expr<V>) {
      return evaluate(m, f, k.value());
    } else {
      auto v_index      = nda::make_regular(nda::transpose(m.units_inv()) * nda::basic_array_view{k});
      auto g            = [&f](long x, long y, long z) { return f(typename brzone::index_t{x, y, z}); };
      auto [d0, d1, d2] = m.dims();
      return evaluate(std::tuple{detail::brzone1d{d0}, detail::brzone1d{d1}, detail::brzone1d{d2}}, g, v_index[0], v_index[1], v_index[2]);
    }
  }



  template <typename T>
  concept BzMeshPoint = is_k_expr<std::decay_t<T>> or std::is_same_v<std::decay_t<T>, brzone::mesh_point_t>;

  template <BzMeshPoint L> k_expr_unary<'-', L> operator-(L &&l) { return {std::forward<L>(l)}; }

  template <BzMeshPoint L, BzMeshPoint R> k_expr<'+', L, R> operator+(L &&l, R &&r) { return {std::forward<L>(l), std::forward<R>(r)}; }

  template <BzMeshPoint L, BzMeshPoint R> k_expr<'-', L, R> operator-(L &&l, R &&r) { return {std::forward<L>(l), std::forward<R>(r)}; }

  template <std::integral Int, BzMeshPoint R> k_expr<'*', long, R> operator*(Int l, R &&r) { return {static_cast<long>(l), std::forward<R>(r)}; }

 
} // namespace triqs::mesh