lattice.h

Go to the documentation of this file.

/*
 *
 *    Copyright (c) 2015-2021,2023-2024
 *      SMASH Team
 *
 *    GNU General Public License (GPLv3 or later)
 *
 */
 
#ifndef SRC_INCLUDE_SMASH_LATTICE_H_
#define SRC_INCLUDE_SMASH_LATTICE_H_
 
#include <array>
#include <cstring>
#include <functional>
#include <optional>
#include <utility>
#include <vector>
 
#include "forwarddeclarations.h"
#include "fourvector.h"
#include "logging.h"
#include "numeric_cast.h"
#include "numerics.h"
 
namespace smash {
static constexpr int LLattice = LogArea::Lattice::id;
 
enum class LatticeUpdate {
  AtOutput = 0,
  EveryTimestep = 1,
  EveryFixedInterval = 2,
};
 
template <typename T>
class RectangularLattice {
 public:
  RectangularLattice(const std::array<double, 3>& l,
                     const std::array<int, 3>& n,
                     const std::array<double, 3>& orig, bool per,
                     const LatticeUpdate upd)
      : lattice_sizes_(l),
        n_cells_(n),
        cell_sizes_{l[0] / n[0], l[1] / n[1], l[2] / n[2]},
        cell_volume_{cell_sizes_[0] * cell_sizes_[1] * cell_sizes_[2]},
        origin_(orig),
        periodic_(per),
        when_update_(upd) {
    lattice_.resize(n_cells_[0] * n_cells_[1] * n_cells_[2]);
    logg[LLattice].debug(
        "Rectangular lattice created: sizes[fm] = (", lattice_sizes_[0], ",",
        lattice_sizes_[1], ",", lattice_sizes_[2], "), dims = (", n_cells_[0],
        ",", n_cells_[1], ",", n_cells_[2], "), origin = (", origin_[0], ",",
        origin_[1], ",", origin_[2], "), periodic: ", periodic_);
    if (n_cells_[0] < 1 || n_cells_[1] < 1 || n_cells_[2] < 1 ||
        lattice_sizes_[0] < 0.0 || lattice_sizes_[1] < 0.0 ||
        lattice_sizes_[2] < 0.0) {
      throw std::invalid_argument(
          "Lattice sizes should be positive, "
          "number of lattice cells should be > 0.");
    }
  }
 
  RectangularLattice(RectangularLattice<T> const& rl)
      : lattice_(rl.lattice_),
        lattice_sizes_(rl.lattice_sizes_),
        n_cells_(rl.n_cells_),
        cell_sizes_(rl.cell_sizes_),
        cell_volume_(rl.cell_volume_),
        origin_(rl.origin_),
        periodic_(rl.periodic_),
        when_update_(rl.when_update_) {}
 
  void reset() { std::fill(lattice_.begin(), lattice_.end(), T()); }
 
  inline bool out_of_bounds(int ix, int iy, int iz) const {
    // clang-format off
    return !periodic_ &&
        (ix < 0 || ix >= n_cells_[0] ||
         iy < 0 || iy >= n_cells_[1] ||
         iz < 0 || iz >= n_cells_[2]);
    // clang-format on
  }
 
  inline ThreeVector cell_center(int ix, int iy, int iz) const {
    return ThreeVector(origin_[0] + cell_sizes_[0] * (ix + 0.5),
                       origin_[1] + cell_sizes_[1] * (iy + 0.5),
                       origin_[2] + cell_sizes_[2] * (iz + 0.5));
  }
 
  inline ThreeVector cell_center(int index) const {
    const int ix = index % n_cells_[0];
    index = index / n_cells_[0];
    const int iy = index % n_cells_[1];
    const int iz = index / n_cells_[1];
    return cell_center(ix, iy, iz);
  }
 
  const std::array<double, 3>& lattice_sizes() const { return lattice_sizes_; }
 
  const std::array<int, 3>& n_cells() const { return n_cells_; }
 
  const std::array<double, 3>& cell_sizes() const { return cell_sizes_; }
 
  const std::array<double, 3>& origin() const { return origin_; }
 
  bool periodic() const { return periodic_; }
 
  LatticeUpdate when_update() const { return when_update_; }
 
  using iterator = typename std::vector<T>::iterator;
  using const_iterator = typename std::vector<T>::const_iterator;
  iterator begin() { return lattice_.begin(); }
  const_iterator begin() const { return lattice_.begin(); }
  iterator end() { return lattice_.end(); }
  const_iterator end() const { return lattice_.end(); }
  T& operator[](std::size_t i) { return lattice_[i]; }
  const T& operator[](std::size_t i) const { return lattice_[i]; }
  std::size_t size() const { return lattice_.size(); }
 
  void assign_value(int lattice_index, T value) {
    lattice_[lattice_index] = value;
  }
 
  int index1d(int ix, int iy, int iz) {
    assert(ix < n_cells_[0]);
    assert(iy < n_cells_[1]);
    assert(iz < n_cells_[2]);
    return (ix + n_cells_[0] * (iy + iz * n_cells_[1]));
  }
  int index_left(int ix, int iy, int iz) {
    if (unlikely(ix == 0)) {
      return periodic_ ? index1d(n_cells_[0] - 1, iy, iz) : index1d(ix, iy, iz);
    } else {
      return index1d(ix - 1, iy, iz);
    }
  }
  int index_right(int ix, int iy, int iz) {
    if (unlikely(ix == (n_cells_[0] - 1))) {
      return periodic_ ? index1d(0, iy, iz) : index1d(ix, iy, iz);
    } else {
      return index1d(ix + 1, iy, iz);
    }
  }
  int index_down(int ix, int iy, int iz) {
    if (unlikely(iy == 0)) {
      return periodic_ ? index1d(ix, n_cells_[1] - 1, iz) : index1d(ix, iy, iz);
    } else {
      return index1d(ix, iy - 1, iz);
    }
  }
  int index_up(int ix, int iy, int iz) {
    if (unlikely(iy == (n_cells_[1] - 1))) {
      return periodic_ ? index1d(ix, 0, iz) : index1d(ix, iy, iz);
    } else {
      return index1d(ix, iy + 1, iz);
    }
  }
  int index_backward(int ix, int iy, int iz) {
    if (unlikely(iz == 0)) {
      return periodic_ ? index1d(ix, iy, n_cells_[2] - 1) : index1d(ix, iy, iz);
    } else {
      return index1d(ix, iy, iz - 1);
    }
  }
  int index_forward(int ix, int iy, int iz) {
    if (unlikely(iz == (n_cells_[2] - 1))) {
      return periodic_ ? index1d(ix, iy, 0) : index1d(ix, iy, iz);
    } else {
      return index1d(ix, iy, iz + 1);
    }
  }
 
  void compute_gradient_lattice(
      RectangularLattice<ThreeVector>& grad_lat) const {
    if (n_cells_[0] < 2 || n_cells_[1] < 2 || n_cells_[2] < 2) {
      // Gradient calculation is impossible
      throw std::runtime_error(
          "Lattice is too small for gradient calculation"
          " (should be at least 2x2x2)");
    }
    if (!identical_to_lattice(&grad_lat)) {
      // Lattice for gradient should have identical origin/dims/periodicity
      throw std::invalid_argument(
          "Lattice for gradient should have the"
          " same origin/dims/periodicity as the original one.");
    }
    const double inv_2dx = 0.5 / cell_sizes_[0];
    const double inv_2dy = 0.5 / cell_sizes_[1];
    const double inv_2dz = 0.5 / cell_sizes_[2];
    const int dix = 1;
    const int diy = n_cells_[0];
    const int diz = n_cells_[0] * n_cells_[1];
    const int d = diz * n_cells_[2];
 
    for (int iz = 0; iz < n_cells_[2]; iz++) {
      const int z_offset = diz * iz;
      for (int iy = 0; iy < n_cells_[1]; iy++) {
        const int y_offset = diy * iy + z_offset;
        for (int ix = 0; ix < n_cells_[0]; ix++) {
          const int index = ix + y_offset;
          if (unlikely(ix == 0)) {
            (grad_lat)[index].set_x1(
                periodic_
                    ? (lattice_[index + dix] - lattice_[index + diy - dix]) *
                          inv_2dx
                    : (lattice_[index + dix] - lattice_[index]) * 2.0 *
                          inv_2dx);
          } else if (unlikely(ix == n_cells_[0] - 1)) {
            (grad_lat)[index].set_x1(
                periodic_
                    ? (lattice_[index - diy + dix] - lattice_[index - dix]) *
                          inv_2dx
                    : (lattice_[index] - lattice_[index - dix]) * 2.0 *
                          inv_2dx);
          } else {
            (grad_lat)[index].set_x1(
                (lattice_[index + dix] - lattice_[index - dix]) * inv_2dx);
          }
 
          if (unlikely(iy == 0)) {
            (grad_lat)[index].set_x2(
                periodic_
                    ? (lattice_[index + diy] - lattice_[index + diz - diy]) *
                          inv_2dy
                    : (lattice_[index + diy] - lattice_[index]) * 2.0 *
                          inv_2dy);
          } else if (unlikely(iy == n_cells_[1] - 1)) {
            (grad_lat)[index].set_x2(
                periodic_
                    ? (lattice_[index - diz + diy] - lattice_[index - diy]) *
                          inv_2dy
                    : (lattice_[index] - lattice_[index - diy]) * 2.0 *
                          inv_2dy);
          } else {
            (grad_lat)[index].set_x2(
                (lattice_[index + diy] - lattice_[index - diy]) * inv_2dy);
          }
 
          if (unlikely(iz == 0)) {
            (grad_lat)[index].set_x3(
                periodic_
                    ? (lattice_[index + diz] - lattice_[index + d - diz]) *
                          inv_2dz
                    : (lattice_[index + diz] - lattice_[index]) * 2.0 *
                          inv_2dz);
          } else if (unlikely(iz == n_cells_[2] - 1)) {
            (grad_lat)[index].set_x3(
                periodic_
                    ? (lattice_[index - d + diz] - lattice_[index - diz]) *
                          inv_2dz
                    : (lattice_[index] - lattice_[index - diz]) * 2.0 *
                          inv_2dz);
          } else {
            (grad_lat)[index].set_x3(
                (lattice_[index + diz] - lattice_[index - diz]) * inv_2dz);
          }
        }
      }
    }
  }
 
  void compute_four_gradient_lattice(
      RectangularLattice<FourVector>& old_lat, double time_step,
      RectangularLattice<std::array<FourVector, 4>>& grad_lat) const {
    if (n_cells_[0] < 2 || n_cells_[1] < 2 || n_cells_[2] < 2) {
      // Gradient calculation is impossible
      throw std::runtime_error(
          "Lattice is too small for gradient calculation"
          " (should be at least 2x2x2)");
    }
    if (!identical_to_lattice(&grad_lat)) {
      // Lattice for gradient should have identical origin/dims/periodicity
      throw std::invalid_argument(
          "Lattice for gradient should have the"
          " same origin/dims/periodicity as the original one.");
    }
    const double inv_2dx = 0.5 / cell_sizes_[0];
    const double inv_2dy = 0.5 / cell_sizes_[1];
    const double inv_2dz = 0.5 / cell_sizes_[2];
    const int dix = 1;
    const int diy = n_cells_[0];
    const int diz = n_cells_[0] * n_cells_[1];
    const int d = diz * n_cells_[2];
 
    for (int iz = 0; iz < n_cells_[2]; iz++) {
      const int z_offset = diz * iz;
      for (int iy = 0; iy < n_cells_[1]; iy++) {
        const int y_offset = diy * iy + z_offset;
        for (int ix = 0; ix < n_cells_[0]; ix++) {
          const int index = ix + y_offset;
 
          // auxiliary vectors used for the calculation of gradients
          FourVector grad_t_jmu(0.0, 0.0, 0.0, 0.0);
          FourVector grad_x_jmu(0.0, 0.0, 0.0, 0.0);
          FourVector grad_y_jmu(0.0, 0.0, 0.0, 0.0);
          FourVector grad_z_jmu(0.0, 0.0, 0.0, 0.0);
          // t direction
          grad_t_jmu = (lattice_[index] - (old_lat)[index]) * (1.0 / time_step);
          // x direction
          if (unlikely(ix == 0)) {
            grad_x_jmu =
                periodic_
                    ? (lattice_[index + dix] - lattice_[index + diy - dix]) *
                          inv_2dx
                    : (lattice_[index + dix] - lattice_[index]) * 2.0 * inv_2dx;
          } else if (unlikely(ix == n_cells_[0] - 1)) {
            grad_x_jmu =
                periodic_
                    ? (lattice_[index - diy + dix] - lattice_[index - dix]) *
                          inv_2dx
                    : (lattice_[index] - lattice_[index - dix]) * 2.0 * inv_2dx;
          } else {
            grad_x_jmu =
                (lattice_[index + dix] - lattice_[index - dix]) * inv_2dx;
          }
          // y direction
          if (unlikely(iy == 0)) {
            grad_y_jmu =
                periodic_
                    ? (lattice_[index + diy] - lattice_[index + diz - diy]) *
                          inv_2dy
                    : (lattice_[index + diy] - lattice_[index]) * 2.0 * inv_2dy;
          } else if (unlikely(iy == n_cells_[1] - 1)) {
            grad_y_jmu =
                periodic_
                    ? (lattice_[index - diz + diy] - lattice_[index - diy]) *
                          inv_2dy
                    : (lattice_[index] - lattice_[index - diy]) * 2.0 * inv_2dy;
          } else {
            grad_y_jmu =
                (lattice_[index + diy] - lattice_[index - diy]) * inv_2dy;
          }
          // z direction
          if (unlikely(iz == 0)) {
            grad_z_jmu =
                periodic_
                    ? (lattice_[index + diz] - lattice_[index + d - diz]) *
                          inv_2dz
                    : (lattice_[index + diz] - lattice_[index]) * 2.0 * inv_2dz;
          } else if (unlikely(iz == n_cells_[2] - 1)) {
            grad_z_jmu =
                periodic_
                    ? (lattice_[index - d + diz] - lattice_[index - diz]) *
                          inv_2dz
                    : (lattice_[index] - lattice_[index - diz]) * 2.0 * inv_2dz;
          } else {
            grad_z_jmu =
                (lattice_[index + diz] - lattice_[index - diz]) * inv_2dz;
          }
          // fill
          (grad_lat)[index][0] = grad_t_jmu;
          (grad_lat)[index][1] = grad_x_jmu;
          (grad_lat)[index][2] = grad_y_jmu;
          (grad_lat)[index][3] = grad_z_jmu;
        }
      }
    }
  }
 
  const T& node(int ix, int iy, int iz) const {
    return periodic_
               ? lattice_[positive_modulo(ix, n_cells_[0]) +
                          n_cells_[0] *
                              (positive_modulo(iy, n_cells_[1]) +
                               n_cells_[1] * positive_modulo(iz, n_cells_[2]))]
               : lattice_[ix + n_cells_[0] * (iy + n_cells_[1] * iz)];
  }
 
  bool value_at(const ThreeVector& r, T& value) const {
    const int ix =
        numeric_cast<int>(std::floor((r.x1() - origin_[0]) / cell_sizes_[0]));
    const int iy =
        numeric_cast<int>(std::floor((r.x2() - origin_[1]) / cell_sizes_[1]));
    const int iz =
        numeric_cast<int>(std::floor((r.x3() - origin_[2]) / cell_sizes_[2]));
    if (out_of_bounds(ix, iy, iz)) {
      value = T();
      return false;
    } else {
      value = node(ix, iy, iz);
      return true;
    }
  }
 
  template <typename F>
  void iterate_sublattice(const std::array<int, 3>& lower_bounds,
                          const std::array<int, 3>& upper_bounds, F&& func) {
    logg[LLattice].debug(
        "Iterating sublattice with lower bound index (", lower_bounds[0], ",",
        lower_bounds[1], ",", lower_bounds[2], "), upper bound index (",
        upper_bounds[0], ",", upper_bounds[1], ",", upper_bounds[2], ")");
 
    if (periodic_) {
      for (int iz = lower_bounds[2]; iz < upper_bounds[2]; iz++) {
        const int z_offset = positive_modulo(iz, n_cells_[2]) * n_cells_[1];
        for (int iy = lower_bounds[1]; iy < upper_bounds[1]; iy++) {
          const int y_offset =
              n_cells_[0] * (positive_modulo(iy, n_cells_[1]) + z_offset);
          for (int ix = lower_bounds[0]; ix < upper_bounds[0]; ix++) {
            const int index = positive_modulo(ix, n_cells_[0]) + y_offset;
            func(lattice_[index], ix, iy, iz);
          }
        }
      }
    } else {
      for (int iz = lower_bounds[2]; iz < upper_bounds[2]; iz++) {
        const int z_offset = iz * n_cells_[1];
        for (int iy = lower_bounds[1]; iy < upper_bounds[1]; iy++) {
          const int y_offset = n_cells_[0] * (iy + z_offset);
          for (int ix = lower_bounds[0]; ix < upper_bounds[0]; ix++) {
            func(lattice_[ix + y_offset], ix, iy, iz);
          }
        }
      }
    }
  }
 
  template <typename F>
  void iterate_in_cube(const ThreeVector& point, const double r_cut, F&& func) {
    std::array<int, 3> l_bounds, u_bounds;
 
    /* Array holds value at the cell center: r_center = r_0 + (i+0.5)cell_size,
     * where i is index in any direction. Therefore we want cells with condition
     * (r-r_cut)/csize - 0.5 < i < (r+r_cut)/csize - 0.5, r = r_center - r_0 */
    for (int i = 0; i < 3; i++) {
      l_bounds[i] = numeric_cast<int>(
          std::ceil((point[i] - origin_[i] - r_cut) / cell_sizes_[i] - 0.5));
      u_bounds[i] = numeric_cast<int>(
          std::ceil((point[i] - origin_[i] + r_cut) / cell_sizes_[i] - 0.5));
    }
 
    if (!periodic_) {
      for (int i = 0; i < 3; i++) {
        if (l_bounds[i] < 0) {
          l_bounds[i] = 0;
        }
        if (u_bounds[i] > n_cells_[i]) {
          u_bounds[i] = n_cells_[i];
        }
        if (l_bounds[i] > n_cells_[i] || u_bounds[i] < 0) {
          return;
        }
      }
    }
    iterate_sublattice(l_bounds, u_bounds, std::forward<F>(func));
  }
 
  template <typename F>
  void integrate_volume(F& integral,
                        F (*integrand)(ThreeVector, T&, ThreeVector),
                        const double rcut, const ThreeVector& point) {
    iterate_in_rectangle(
        point, {rcut, rcut, rcut},
        [&point, &integral, &integrand, this](T value, int ix, int iy, int iz) {
          ThreeVector pos = this->cell_center(ix, iy, iz);
          integral += integrand(pos, value, point) * this->cell_volume_;
        });
  }
  template <typename F>
  void iterate_in_rectangle(const ThreeVector& point,
                            const std::array<double, 3>& rectangle, F&& func) {
    std::array<int, 3> l_bounds, u_bounds;
 
    /* Array holds value at the cell center: r_center = r_0 + (i+0.5)cell_size,
     * where i is index in any direction. Therefore we want cells with condition
     * (r[i]-rectangle[i])/csize - 0.5 < i < (r[i]+rectangle[i])/csize - 0.5,
     * r[i] = r_center[i] - r_0[i]
     */
    for (int i = 0; i < 3; i++) {
      l_bounds[i] = numeric_cast<int>(std::ceil(
          (point[i] - origin_[i] - rectangle[i]) / cell_sizes_[i] - 0.5));
      u_bounds[i] = numeric_cast<int>(std::ceil(
          (point[i] - origin_[i] + rectangle[i]) / cell_sizes_[i] - 0.5));
    }
 
    if (!periodic_) {
      for (int i = 0; i < 3; i++) {
        if (l_bounds[i] < 0) {
          l_bounds[i] = 0;
        }
        if (u_bounds[i] > n_cells_[i]) {
          u_bounds[i] = n_cells_[i];
        }
        if (l_bounds[i] > n_cells_[i] || u_bounds[i] < 0) {
          return;
        }
      }
    }
    iterate_sublattice(l_bounds, u_bounds, std::forward<F>(func));
  }
 
  template <typename F>
  void iterate_nearest_neighbors(const ThreeVector& point, F&& func) {
    // get the 3D indices of the cell containing the given point
    const int ix = numeric_cast<int>(
        std::floor((point.x1() - origin_[0]) / cell_sizes_[0]));
    const int iy = numeric_cast<int>(
        std::floor((point.x2() - origin_[1]) / cell_sizes_[1]));
    const int iz = numeric_cast<int>(
        std::floor((point.x3() - origin_[2]) / cell_sizes_[2]));
 
    logg[LLattice].debug(
        "Iterating over nearest neighbors of the cell at ix = ", ix,
        ", iy = ", iy, ", iz = ", iz);
 
    // determine the 1D index of the center cell
    const int i = index1d(ix, iy, iz);
    func(lattice_[i], i, i);
 
    // determine the indeces of nearby cells, perform function on them
    int ileft = index_left(ix, iy, iz);
    func(lattice_[ileft], ileft, i);
 
    int iright = index_right(ix, iy, iz);
    func(lattice_[iright], iright, i);
 
    int idown = index_down(ix, iy, iz);
    func(lattice_[idown], idown, i);
 
    int iup = index_up(ix, iy, iz);
    func(lattice_[iup], iup, i);
 
    int ibackward = index_backward(ix, iy, iz);
    func(lattice_[ibackward], ibackward, i);
 
    int iforward = index_forward(ix, iy, iz);
    func(lattice_[iforward], iforward, i);
  }
 
  template <typename L>
  bool identical_to_lattice(const L* lat) const {
    return n_cells_[0] == lat->n_cells()[0] &&
           n_cells_[1] == lat->n_cells()[1] &&
           n_cells_[2] == lat->n_cells()[2] &&
           std::abs(lattice_sizes_[0] - lat->lattice_sizes()[0]) <
               really_small &&
           std::abs(lattice_sizes_[1] - lat->lattice_sizes()[1]) <
               really_small &&
           std::abs(lattice_sizes_[2] - lat->lattice_sizes()[2]) <
               really_small &&
           std::abs(origin_[0] - lat->origin()[0]) < really_small &&
           std::abs(origin_[1] - lat->origin()[1]) < really_small &&
           std::abs(origin_[2] - lat->origin()[2]) < really_small &&
           periodic_ == lat->periodic();
  }
 
  void reset_and_resize(std::optional<std::array<double, 3>> new_length,
                        std::optional<std::array<double, 3>> new_origin,
                        std::optional<std::array<int, 3>> new_cells) {
    if (!new_length && !new_origin && !new_cells) {
      throw std::invalid_argument(
          "RectangularLattice::reset_and_resize called "
          "without arguments, lattice was not changed.");
    }
    reset();
    if (new_length)
      lattice_sizes_ = *new_length;
    if (new_origin)
      origin_ = *new_origin;
    if (new_cells)
      n_cells_ = *new_cells;
    cell_sizes_ = {lattice_sizes_[0] / n_cells_[0],
                   lattice_sizes_[1] / n_cells_[1],
                   lattice_sizes_[2] / n_cells_[2]};
    cell_volume_ = cell_sizes_[0] * cell_sizes_[1] * cell_sizes_[2];
  }
 
 protected:
  std::vector<T> lattice_;
  std::array<double, 3> lattice_sizes_;
  std::array<int, 3> n_cells_;
  std::array<double, 3> cell_sizes_;
  double cell_volume_;
  std::array<double, 3> origin_;
  const bool periodic_;
  const LatticeUpdate when_update_;
 
 private:
  inline int positive_modulo(int i, int n) const {
    /* (i % n + n) % n would be correct, but slow.
     * Instead I rely on the fact that i should never go too far
     * in negative region and replace i%n + n by i + 256 * n = i + (n << 8) */
    // FIXME: This should use asserts, also checking for under- or overflows.
    return (i + (n << 8)) % n;
  }
};
 
}  // namespace smash
 
#endif  // SRC_INCLUDE_SMASH_LATTICE_H_