LBWalberlaImpl.hpp

Go to the documentation of this file.

/*
 * Copyright (C) 2019-2023 The ESPResSo project
 *
 * This file is part of ESPResSo.
 *
 * ESPResSo 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.
 *
 * ESPResSo 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 should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
 
#pragma once
 
/**
 * @file
 * @ref walberla::LBWalberlaImpl implements the interface of the LB
 * waLBerla bridge using sweeps generated by lbmpy
 * (see <tt>maintainer/walberla_kernels</tt>).
 */
 
#include <blockforest/Initialization.h>
#include <blockforest/StructuredBlockForest.h>
#include <blockforest/communication/UniformBufferedScheme.h>
#include <domain_decomposition/BlockDataID.h>
#include <domain_decomposition/IBlock.h>
#include <field/AddToStorage.h>
#include <field/GhostLayerField.h>
#include <field/communication/PackInfo.h>
#include <field/vtk/FlagFieldCellFilter.h>
#include <field/vtk/VTKWriter.h>
#include <stencil/D3Q19.h>
#include <stencil/D3Q27.h>
#if defined(__CUDACC__)
#include <gpu/AddGPUFieldToStorage.h>
#include <gpu/communication/MemcpyPackInfo.h>
#include <gpu/communication/UniformGPUScheme.h>
#endif
 
#include "../BoundaryHandling.hpp"
#include "../BoundaryPackInfo.hpp"
#include "../utils/boundary.hpp"
#include "../utils/types_conversion.hpp"
#include "InterpolateAndShiftAtBoundary.hpp"
#include "ResetForce.hpp"
#include "lb_kernels.hpp"
#if defined(__CUDACC__)
#include "lb_kernels.cuh"
#endif
 
#include <walberla_bridge/Architecture.hpp>
#include <walberla_bridge/BlockAndCell.hpp>
#include <walberla_bridge/LatticeWalberla.hpp>
#include <walberla_bridge/lattice_boltzmann/LBWalberlaBase.hpp>
#include <walberla_bridge/lattice_boltzmann/LeesEdwardsPack.hpp>
 
#include <utils/Vector.hpp>
#include <utils/interpolation/bspline_3d.hpp>
#include <utils/math/make_lin_space.hpp>
 
#include <array>
#include <cmath>
#include <cstddef>
#include <functional>
#include <initializer_list>
#include <limits>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>
 
namespace walberla {
 
/** @brief Class that runs and controls the LB on waLBerla. */
template <typename FloatType, lbmpy::Arch Architecture>
class LBWalberlaImpl : public LBWalberlaBase {
protected:
  using CollisionModelLeesEdwards =
      typename detail::KernelTrait<FloatType,
                                   Architecture>::CollisionModelLeesEdwards;
  using CollisionModelThermalized =
      typename detail::KernelTrait<FloatType,
                                   Architecture>::CollisionModelThermalized;
  using StreamSweep =
      typename detail::KernelTrait<FloatType, Architecture>::StreamSweep;
  using InitialPDFsSetter =
      typename detail::KernelTrait<FloatType, Architecture>::InitialPDFsSetter;
  using BoundaryModel =
      BoundaryHandling<Vector3<FloatType>,
                       typename detail::BoundaryHandlingTrait<
                           FloatType, Architecture>::Dynamic_UBB>;
  using CollisionModel =
      std::variant<CollisionModelThermalized, CollisionModelLeesEdwards>;
 
public:
  /** @brief Stencil for collision and streaming operations. */
  using Stencil = stencil::D3Q19;
  /** @brief Stencil for ghost communication (includes domain corners). */
  using StencilFull = stencil::D3Q27;
  /** @brief Lattice model (e.g. blockforest). */
  using Lattice_T = LatticeWalberla::Lattice_T;
 
protected:
  template <typename FT, lbmpy::Arch AT = lbmpy::Arch::CPU> struct FieldTrait {
    using PdfField = field::GhostLayerField<FT, Stencil::Size>;
    using VectorField = field::GhostLayerField<FT, uint_t{3u}>;
    template <class Field>
    using PackInfo = field::communication::PackInfo<Field>;
    template <class Stencil>
    using RegularCommScheme =
        blockforest::communication::UniformBufferedScheme<Stencil>;
    template <class Stencil>
    using BoundaryCommScheme =
        blockforest::communication::UniformBufferedScheme<Stencil>;
  };
 
#if defined(__CUDACC__)
  template <typename FT> struct FieldTrait<FT, lbmpy::Arch::GPU> {
    using PdfField = gpu::GPUField<FT>;
    using VectorField = gpu::GPUField<FT>;
    template <class Field>
    using PackInfo = gpu::communication::MemcpyPackInfo<Field>;
    template <class Stencil>
    using RegularCommScheme = gpu::communication::UniformGPUScheme<Stencil>;
    template <class Stencil>
    using BoundaryCommScheme =
        blockforest::communication::UniformBufferedScheme<Stencil>;
  };
#endif
 
  // "underlying" field types (`GPUField` has no f-size info at compile time)
  using _PdfField = typename FieldTrait<FloatType>::PdfField;
  using _VectorField = typename FieldTrait<FloatType>::VectorField;
 
public:
  using PdfField = typename FieldTrait<FloatType, Architecture>::PdfField;
  using VectorField = typename FieldTrait<FloatType, Architecture>::VectorField;
  using FlagField = typename BoundaryModel::FlagField;
#if defined(__CUDACC__)
  using GPUField = gpu::GPUField<FloatType>;
#endif
 
public:
  template <typename T> FloatType FloatType_c(T t) const {
    return numeric_cast<FloatType>(t);
  }
 
  [[nodiscard]] std::size_t stencil_size() const noexcept override {
    return static_cast<std::size_t>(Stencil::Size);
  }
 
  [[nodiscard]] virtual bool is_double_precision() const noexcept override {
    return std::is_same_v<FloatType, double>;
  }
 
private:
  class CollideSweepVisitor {
  public:
    void operator()(CollisionModelThermalized &cm, IBlock *b) { cm(b); }
 
    void operator()(CollisionModelLeesEdwards &cm, IBlock *b) {
      cm.v_s_ = static_cast<decltype(cm.v_s_)>(
          m_lees_edwards_callbacks->get_shear_velocity());
      cm(b);
    }
 
    CollideSweepVisitor() = default;
    explicit CollideSweepVisitor(std::shared_ptr<LeesEdwardsPack> callbacks) {
      m_lees_edwards_callbacks = std::move(callbacks);
    }
 
  private:
    std::shared_ptr<LeesEdwardsPack> m_lees_edwards_callbacks{};
  };
  CollideSweepVisitor m_run_collide_sweep{};
 
  FloatType shear_mode_relaxation_rate() const {
    return FloatType{2} / (FloatType{6} * m_viscosity + FloatType{1});
  }
 
  FloatType odd_mode_relaxation_rate(
      FloatType shear_relaxation,
      FloatType magic_number = FloatType{3} / FloatType{16}) const {
    return (FloatType{4} - FloatType{2} * shear_relaxation) /
           (FloatType{4} * magic_number * shear_relaxation + FloatType{2} -
            shear_relaxation);
  }
 
  void reset_boundary_handling() {
    auto const &blocks = get_lattice().get_blocks();
    m_boundary = std::make_shared<BoundaryModel>(blocks, m_pdf_field_id,
                                                 m_flag_field_id);
  }
 
  FloatType pressure_tensor_correction_factor() const {
    return m_viscosity / (m_viscosity + FloatType{1} / FloatType{6});
  }
 
  void pressure_tensor_correction(Matrix3<FloatType> &tensor) const {
    auto const revert_factor = pressure_tensor_correction_factor();
    for (auto const i : {1u, 2u, 3u, 5u, 6u, 7u}) {
      tensor[i] *= revert_factor;
    }
  }
 
  class interpolation_illegal_access : public std::runtime_error {
  public:
    explicit interpolation_illegal_access(std::string const &field,
                                          Utils::Vector3d const &pos,
                                          std::array<int, 3> const &node,
                                          double weight)
        : std::runtime_error("Access to LB " + field + " field failed") {
      std::cerr << "pos [" << pos << "], node [" << Utils::Vector3i(node)
                << "], weight " << weight << "\n";
    }
  };
 
  class vtk_runtime_error : public std::runtime_error {
  public:
    explicit vtk_runtime_error(std::string const &vtk_uid,
                               std::string const &reason)
        : std::runtime_error("VTKOutput object '" + vtk_uid + "' " + reason) {}
  };
 
protected:
  // Member variables
  FloatType m_viscosity; /// kinematic viscosity
  FloatType m_density;
  FloatType m_kT;
 
  // Block data access handles
  BlockDataID m_pdf_field_id;
  BlockDataID m_pdf_tmp_field_id;
  BlockDataID m_flag_field_id;
 
  BlockDataID m_last_applied_force_field_id;
  BlockDataID m_force_to_be_applied_id;
 
  BlockDataID m_velocity_field_id;
  BlockDataID m_vec_tmp_field_id;
 
  /** Flag for boundary cells. */
  FlagUID const Boundary_flag{"boundary"};
 
  /**
   * @brief Full communicator.
   * We use the D3Q27 directions to update cells along the diagonals during
   * a full ghost communication. This is needed to properly update the corners
   * of the ghost layer when setting cell velocities or populations.
   */
  using RegularFullCommunicator =
      typename FieldTrait<FloatType, Architecture>::template RegularCommScheme<
          typename stencil::D3Q27>;
  using BoundaryFullCommunicator =
      typename FieldTrait<FloatType, Architecture>::template BoundaryCommScheme<
          typename stencil::D3Q27>;
  /**
   * @brief Regular communicator.
   * We use the same directions as the stencil during integration.
   */
  using PDFStreamingCommunicator =
      typename FieldTrait<FloatType,
                          Architecture>::template RegularCommScheme<Stencil>;
  template <class Field>
  using PackInfo =
      typename FieldTrait<FloatType, Architecture>::template PackInfo<Field>;
 
  // communicators
  std::shared_ptr<BoundaryFullCommunicator> m_boundary_communicator;
  std::shared_ptr<RegularFullCommunicator> m_pdf_full_communicator;
  std::shared_ptr<PDFStreamingCommunicator> m_pdf_streaming_communicator;
 
  // ResetForce sweep + external force handling
  std::shared_ptr<ResetForce<PdfField, VectorField>> m_reset_force;
 
  // Stream sweep
  std::shared_ptr<StreamSweep> m_stream;
 
  // Lees Edwards boundary interpolation
  std::shared_ptr<LeesEdwardsPack> m_lees_edwards_callbacks;
  std::shared_ptr<InterpolateAndShiftAtBoundary<_PdfField, FloatType>>
      m_lees_edwards_pdf_interpol_sweep;
  std::shared_ptr<InterpolateAndShiftAtBoundary<_VectorField, FloatType>>
      m_lees_edwards_vel_interpol_sweep;
  std::shared_ptr<InterpolateAndShiftAtBoundary<_VectorField, FloatType>>
      m_lees_edwards_last_applied_force_interpol_sweep;
 
  // Collision sweep
  std::shared_ptr<CollisionModel> m_collision_model;
 
  // boundaries
  std::shared_ptr<BoundaryModel> m_boundary;
 
  // lattice
  std::shared_ptr<LatticeWalberla> m_lattice;
 
  [[nodiscard]] std::optional<CellInterval>
  get_interval(Utils::Vector3i const &lower_corner,
               Utils::Vector3i const &upper_corner) const {
    auto const &lattice = get_lattice();
    auto const &cell_min = lower_corner;
    auto const cell_max = upper_corner - Utils::Vector3i::broadcast(1);
    auto const lower_bc = get_block_and_cell(lattice, cell_min, true);
    auto const upper_bc = get_block_and_cell(lattice, cell_max, true);
    if (not lower_bc or not upper_bc) {
      return std::nullopt;
    }
    assert(&(*(lower_bc->block)) == &(*(upper_bc->block)));
    return {CellInterval(lower_bc->cell, upper_bc->cell)};
  }
 
  /**
   * @brief Convenience function to add a field with a custom allocator.
   *
   * When vectorization is off, let waLBerla decide which memory allocator
   * to use. When vectorization is on, the aligned memory allocator is
   * required, otherwise <tt>cpu_vectorize_info["assume_aligned"]</tt> will
   * trigger assertions. That is because for single-precision kernels the
   * waLBerla heuristic in <tt>src/field/allocation/FieldAllocator.h</tt>
   * will fall back to @c StdFieldAlloc, yet @c AllocateAligned is needed
   * for intrinsics to work.
   */
  template <typename Field> auto add_to_storage(std::string const tag) {
    auto const &blocks = m_lattice->get_blocks();
    auto const n_ghost_layers = m_lattice->get_ghost_layers();
    if constexpr (Architecture == lbmpy::Arch::CPU) {
#ifdef ESPRESSO_BUILD_WITH_AVX_KERNELS
#if defined(__AVX512F__)
      constexpr uint_t alignment = 64;
#elif defined(__AVX__)
      constexpr uint_t alignment = 32;
#elif defined(__SSE__)
      constexpr uint_t alignment = 16;
#else
#error "Unsupported arch, check walberla src/field/allocation/FieldAllocator.h"
#endif
      using value_type = typename Field::value_type;
      using Allocator = field::AllocateAligned<value_type, alignment>;
      auto const allocator = std::make_shared<Allocator>();
      auto const empty_set = Set<SUID>::emptySet();
      return field::addToStorage<Field>(
          blocks, tag, field::internal::defaultSize, FloatType{0}, field::fzyx,
          n_ghost_layers, false, {}, empty_set, empty_set, allocator);
#else  // ESPRESSO_BUILD_WITH_AVX_KERNELS
      return field::addToStorage<Field>(blocks, tag, FloatType{0}, field::fzyx,
                                        n_ghost_layers);
#endif // ESPRESSO_BUILD_WITH_AVX_KERNELS
    }
#if defined(__CUDACC__)
    else {
      auto field_id = gpu::addGPUFieldToStorage<GPUField>(
          blocks, tag, Field::F_SIZE, field::fzyx, n_ghost_layers);
      if constexpr (std::is_same_v<Field, _VectorField>) {
        for (auto block = blocks->begin(); block != blocks->end(); ++block) {
          auto field = block->template getData<GPUField>(field_id);
          lbm::accessor::Vector::initialize(field, Vector3<FloatType>{0});
        }
      } else if constexpr (std::is_same_v<Field, _PdfField>) {
        for (auto block = blocks->begin(); block != blocks->end(); ++block) {
          auto field = block->template getData<GPUField>(field_id);
          lbm::accessor::Population::initialize(
              field, std::array<FloatType, Stencil::Size>{});
        }
      }
      return field_id;
    }
#endif
  }
 
public:
  LBWalberlaImpl(std::shared_ptr<LatticeWalberla> lattice, double viscosity,
                 double density)
      : m_viscosity(FloatType_c(viscosity)), m_density(FloatType_c(density)),
        m_kT(FloatType{0}), m_lattice(std::move(lattice)) {
 
    auto const &blocks = m_lattice->get_blocks();
    auto const n_ghost_layers = m_lattice->get_ghost_layers();
    if (n_ghost_layers == 0u)
      throw std::runtime_error("At least one ghost layer must be used");
 
    // Initialize and register fields (must use the "underlying" types)
    m_pdf_field_id = add_to_storage<_PdfField>("pdfs");
    m_pdf_tmp_field_id = add_to_storage<_PdfField>("pdfs_tmp");
    m_last_applied_force_field_id = add_to_storage<_VectorField>("force last");
    m_force_to_be_applied_id = add_to_storage<_VectorField>("force next");
    m_velocity_field_id = add_to_storage<_VectorField>("velocity");
    m_vec_tmp_field_id = add_to_storage<_VectorField>("velocity_tmp");
 
    // Initialize and register pdf field
    auto pdf_setter =
        InitialPDFsSetter(m_force_to_be_applied_id, m_pdf_field_id,
                          m_velocity_field_id, m_density);
    for (auto b = blocks->begin(); b != blocks->end(); ++b) {
      pdf_setter(&*b);
    }
 
    // Initialize and register flag field (fluid/boundary)
    m_flag_field_id = field::addFlagFieldToStorage<FlagField>(
        blocks, "flag field", n_ghost_layers);
    // Initialize boundary sweep
    reset_boundary_handling();
 
    // Set up the communication and register fields
    m_pdf_streaming_communicator =
        std::make_shared<PDFStreamingCommunicator>(blocks);
    m_pdf_streaming_communicator->addPackInfo(
        std::make_shared<PackInfo<PdfField>>(m_pdf_field_id));
    m_pdf_streaming_communicator->addPackInfo(
        std::make_shared<PackInfo<VectorField>>(m_last_applied_force_field_id));
 
    m_pdf_full_communicator = std::make_shared<RegularFullCommunicator>(blocks);
    m_pdf_full_communicator->addPackInfo(
        std::make_shared<PackInfo<PdfField>>(m_pdf_field_id));
    m_pdf_full_communicator->addPackInfo(
        std::make_shared<PackInfo<VectorField>>(m_last_applied_force_field_id));
    m_pdf_full_communicator->addPackInfo(
        std::make_shared<PackInfo<VectorField>>(m_velocity_field_id));
 
    m_boundary_communicator =
        std::make_shared<BoundaryFullCommunicator>(blocks);
    m_boundary_communicator->addPackInfo(
        std::make_shared<field::communication::PackInfo<FlagField>>(
            m_flag_field_id));
    auto boundary_packinfo = std::make_shared<
        field::communication::BoundaryPackInfo<FlagField, BoundaryModel>>(
        m_flag_field_id);
    boundary_packinfo->setup_boundary_handle(m_lattice, m_boundary);
    m_boundary_communicator->addPackInfo(boundary_packinfo);
 
    // Instantiate the sweep responsible for force double buffering and
    // external forces
    m_reset_force = std::make_shared<ResetForce<PdfField, VectorField>>(
        m_last_applied_force_field_id, m_force_to_be_applied_id);
 
    // Prepare LB sweeps
    // Note: For now, combined collide-stream sweeps cannot be used,
    // because the collide-push variant is not supported by lbmpy.
    // The following functors are individual in-place collide and stream steps
    m_stream = std::make_shared<StreamSweep>(
        m_last_applied_force_field_id, m_pdf_field_id, m_velocity_field_id);
  }
 
private:
  void integrate_stream(std::shared_ptr<Lattice_T> const &blocks) {
    for (auto b = blocks->begin(); b != blocks->end(); ++b)
      (*m_stream)(&*b);
  }
 
  void integrate_collide(std::shared_ptr<Lattice_T> const &blocks) {
    auto &cm_variant = *m_collision_model;
    for (auto b = blocks->begin(); b != blocks->end(); ++b)
      std::visit(m_run_collide_sweep, cm_variant, std::variant<IBlock *>(&*b));
    if (auto *cm = std::get_if<CollisionModelThermalized>(&cm_variant)) {
      cm->time_step_++;
    }
  }
 
  auto has_lees_edwards_bc() const {
    return std::holds_alternative<CollisionModelLeesEdwards>(
        *m_collision_model);
  }
 
  void apply_lees_edwards_pdf_interpolation(
      std::shared_ptr<Lattice_T> const &blocks) {
    for (auto b = blocks->begin(); b != blocks->end(); ++b)
      (*m_lees_edwards_pdf_interpol_sweep)(&*b);
  }
 
  void apply_lees_edwards_vel_interpolation_and_shift(
      std::shared_ptr<Lattice_T> const &blocks) {
    for (auto b = blocks->begin(); b != blocks->end(); ++b)
      (*m_lees_edwards_vel_interpol_sweep)(&*b);
  }
 
  void apply_lees_edwards_last_applied_force_interpolation(
      std::shared_ptr<Lattice_T> const &blocks) {
    for (auto b = blocks->begin(); b != blocks->end(); ++b)
      (*m_lees_edwards_last_applied_force_interpol_sweep)(&*b);
  }
 
  void integrate_reset_force(std::shared_ptr<Lattice_T> const &blocks) {
    for (auto b = blocks->begin(); b != blocks->end(); ++b)
      (*m_reset_force)(&*b);
  }
 
  void integrate_boundaries(std::shared_ptr<Lattice_T> const &blocks) {
    for (auto b = blocks->begin(); b != blocks->end(); ++b)
      (*m_boundary)(&*b);
  }
 
  void integrate_push_scheme() {
    auto const &blocks = get_lattice().get_blocks();
    // Reset force fields
    integrate_reset_force(blocks);
    // LB collide
    integrate_collide(blocks);
    m_pdf_streaming_communicator->communicate();
    // Handle boundaries
    integrate_boundaries(blocks);
    // LB stream
    integrate_stream(blocks);
    // Refresh ghost layers
    ghost_communication_pdfs();
  }
 
  void integrate_pull_scheme() {
    auto const &blocks = get_lattice().get_blocks();
    integrate_reset_force(blocks);
    // Handle boundaries
    integrate_boundaries(blocks);
    // LB stream
    integrate_stream(blocks);
    // LB collide
    integrate_collide(blocks);
    // Refresh ghost layers
    ghost_communication_pdfs();
  }
 
protected:
  void integrate_vtk_writers() override {
    for (auto const &it : m_vtk_auto) {
      auto &vtk_handle = it.second;
      if (vtk_handle->enabled) {
        vtk::writeFiles(vtk_handle->ptr)();
        vtk_handle->execution_count++;
      }
    }
  }
 
public:
  void integrate() override {
    if (has_lees_edwards_bc()) {
      integrate_pull_scheme();
    } else {
      integrate_push_scheme();
    }
    // Handle VTK writers
    integrate_vtk_writers();
  }
 
  void ghost_communication() override {
    ghost_communication_boundary();
    ghost_communication_pdfs();
  }
 
  void ghost_communication_boundary() {
    m_boundary_communicator->communicate();
  }
 
  void ghost_communication_pdfs() {
    m_pdf_full_communicator->communicate();
    if (has_lees_edwards_bc()) {
      auto const &blocks = get_lattice().get_blocks();
      apply_lees_edwards_pdf_interpolation(blocks);
      apply_lees_edwards_vel_interpolation_and_shift(blocks);
      apply_lees_edwards_last_applied_force_interpolation(blocks);
    }
  }
 
  void set_collision_model(double kT, unsigned int seed) override {
    auto const omega = shear_mode_relaxation_rate();
    auto const omega_odd = odd_mode_relaxation_rate(omega);
    m_kT = FloatType_c(kT);
    auto obj = CollisionModelThermalized(
        m_last_applied_force_field_id, m_pdf_field_id, uint32_t{0u},
        uint32_t{0u}, uint32_t{0u}, m_kT, omega, omega, omega_odd, omega, seed,
        uint32_t{0u});
    obj.block_offset_generator =
        [this](IBlock *const block, uint32_t &block_offset_0,
               uint32_t &block_offset_1, uint32_t &block_offset_2) {
          auto const &blocks = get_lattice().get_blocks();
          auto const &ci = blocks->getBlockCellBB(*block);
          block_offset_0 = static_cast<uint32_t>(ci.xMin());
          block_offset_1 = static_cast<uint32_t>(ci.yMin());
          block_offset_2 = static_cast<uint32_t>(ci.zMin());
        };
    m_collision_model = std::make_shared<CollisionModel>(std::move(obj));
  }
 
  void set_collision_model(
      std::unique_ptr<LeesEdwardsPack> &&lees_edwards_pack) override {
    assert(m_kT == 0.);
#if defined(__CUDACC__)
    if constexpr (Architecture == lbmpy::Arch::GPU) {
      throw std::runtime_error("Lees-Edwards LB doesn't support GPU yet");
    }
#endif
    auto const shear_direction = lees_edwards_pack->shear_direction;
    auto const shear_plane_normal = lees_edwards_pack->shear_plane_normal;
    auto const shear_vel = FloatType_c(lees_edwards_pack->get_shear_velocity());
    auto const omega = shear_mode_relaxation_rate();
    if (shear_plane_normal != 1u) {
      throw std::domain_error(
          "Lees-Edwards LB only supports shear_plane_normal=\"y\"");
    }
    auto const &lattice = get_lattice();
    auto const n_ghost_layers = lattice.get_ghost_layers();
    auto const blocks = lattice.get_blocks();
    auto const agrid =
        FloatType_c(lattice.get_grid_dimensions()[shear_plane_normal]);
    auto obj = CollisionModelLeesEdwards(
        m_last_applied_force_field_id, m_pdf_field_id, agrid, omega, shear_vel);
    m_collision_model = std::make_shared<CollisionModel>(std::move(obj));
    m_lees_edwards_callbacks = std::move(lees_edwards_pack);
    m_run_collide_sweep = CollideSweepVisitor{m_lees_edwards_callbacks};
    m_lees_edwards_pdf_interpol_sweep =
        std::make_shared<InterpolateAndShiftAtBoundary<_PdfField, FloatType>>(
            blocks, m_pdf_field_id, m_pdf_tmp_field_id, n_ghost_layers,
            shear_direction, shear_plane_normal,
            m_lees_edwards_callbacks->get_pos_offset);
    m_lees_edwards_vel_interpol_sweep = std::make_shared<
        InterpolateAndShiftAtBoundary<_VectorField, FloatType>>(
        blocks, m_velocity_field_id, m_vec_tmp_field_id, n_ghost_layers,
        shear_direction, shear_plane_normal,
        m_lees_edwards_callbacks->get_pos_offset,
        m_lees_edwards_callbacks->get_shear_velocity);
    m_lees_edwards_last_applied_force_interpol_sweep = std::make_shared<
        InterpolateAndShiftAtBoundary<_VectorField, FloatType>>(
        blocks, m_last_applied_force_field_id, m_vec_tmp_field_id,
        n_ghost_layers, shear_direction, shear_plane_normal,
        m_lees_edwards_callbacks->get_pos_offset);
  }
 
  void check_lebc(unsigned int shear_direction,
                  unsigned int shear_plane_normal) const override {
    if (m_lees_edwards_callbacks) {
      if (m_lees_edwards_callbacks->shear_direction != shear_direction or
          m_lees_edwards_callbacks->shear_plane_normal != shear_plane_normal) {
        throw std::runtime_error(
            "MD and LB Lees-Edwards boundary conditions disagree");
      }
    }
  }
 
  void set_viscosity(double viscosity) override {
    m_viscosity = FloatType_c(viscosity);
  }
 
  [[nodiscard]] double get_viscosity() const noexcept override {
    return numeric_cast<double>(m_viscosity);
  }
 
  [[nodiscard]] double get_density() const noexcept override {
    return numeric_cast<double>(m_density);
  }
 
  // Velocity
  std::optional<Utils::Vector3d>
  get_node_velocity(Utils::Vector3i const &node,
                    bool consider_ghosts = false) const override {
    auto const is_boundary = get_node_is_boundary(node, consider_ghosts);
    if (is_boundary)    // is info available locally
      if (*is_boundary) // is the node a boundary
        return get_node_velocity_at_boundary(node, consider_ghosts);
    auto const bc = get_block_and_cell(get_lattice(), node, consider_ghosts);
    if (!bc)
      return std::nullopt;
 
    auto field = bc->block->template uncheckedFastGetData<VectorField>(
        m_velocity_field_id);
    auto const vec = lbm::accessor::Vector::get(field, bc->cell);
    return to_vector3d(vec);
  }
 
  bool set_node_velocity(Utils::Vector3i const &node,
                         Utils::Vector3d const &v) override {
    auto bc = get_block_and_cell(get_lattice(), node, false);
    if (!bc)
      return false;
 
    // We have to set both, the pdf and the stored velocity field
    auto pdf_field = bc->block->template getData<PdfField>(m_pdf_field_id);
    auto vel_field =
        bc->block->template getData<VectorField>(m_velocity_field_id);
    auto force_field =
        bc->block->template getData<VectorField>(m_last_applied_force_field_id);
    auto const vel = to_vector3<FloatType>(v);
    lbm::accessor::Velocity::set(pdf_field, force_field, vel, bc->cell);
    lbm::accessor::Vector::set(vel_field, vel, bc->cell);
 
    return true;
  }
 
  std::vector<double>
  get_slice_velocity(Utils::Vector3i const &lower_corner,
                     Utils::Vector3i const &upper_corner) const override {
    std::vector<double> out;
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto const &block = *(lattice.get_blocks()->begin());
      auto const field =
          block.template getData<VectorField>(m_velocity_field_id);
      auto const local_offset = std::get<0>(lattice.get_local_grid_range());
      auto const lower_cell = ci->min();
      auto const upper_cell = ci->max();
      out.reserve(ci->numCells());
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      for (auto x = lower_cell.x(); x <= upper_cell.x(); ++x) {
        for (auto y = lower_cell.y(); y <= upper_cell.y(); ++y) {
          for (auto z = lower_cell.z(); z <= upper_cell.z(); ++z) {
            auto const node = local_offset + Utils::Vector3i{{x, y, z}};
            if (m_boundary->node_is_boundary(node)) {
              auto const vec = m_boundary->get_node_value_at_boundary(node);
              for (uint_t f = 0u; f < 3u; ++f) {
                out.emplace_back(double_c(vec[f]));
              }
            } else {
              auto const vec = lbm::accessor::Vector::get(field, Cell{x, y, z});
              for (uint_t f = 0u; f < 3u; ++f) {
                out.emplace_back(double_c(vec[f]));
              }
            }
          }
        }
      }
    }
    return out;
  }
 
  void set_slice_velocity(Utils::Vector3i const &lower_corner,
                          Utils::Vector3i const &upper_corner,
                          std::vector<double> const &velocity) override {
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto &block = *(lattice.get_blocks()->begin());
      // We have to set both, the pdf and the stored velocity field
      auto pdf_field = block.template getData<PdfField>(m_pdf_field_id);
      auto vel_field = block.template getData<VectorField>(m_velocity_field_id);
      auto force_field =
          block.template getData<VectorField>(m_last_applied_force_field_id);
      auto const lower_cell = ci->min();
      auto const upper_cell = ci->max();
      auto it = velocity.begin();
      assert(velocity.size() == 3u * ci->numCells());
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      for (auto x = lower_cell.x(); x <= upper_cell.x(); ++x) {
        for (auto y = lower_cell.y(); y <= upper_cell.y(); ++y) {
          for (auto z = lower_cell.z(); z <= upper_cell.z(); ++z) {
            auto const cell = Cell{x, y, z};
            Vector3<FloatType> vec;
            for (uint_t f = 0u; f < 3u; ++f) {
              vec[f] = FloatType_c(*it);
              ++it;
            }
            lbm::accessor::Velocity::set(pdf_field, force_field, vec, cell);
            lbm::accessor::Vector::set(vel_field, vec, cell);
          }
        }
      }
    }
  }
 
  [[nodiscard]] bool is_gpu() const noexcept override {
    return Architecture == lbmpy::Arch::GPU;
  }
 
  void add_forces_at_pos(std::vector<Utils::Vector3d> const &pos,
                         std::vector<Utils::Vector3d> const &forces) override {
    assert(pos.size() == forces.size());
    if (pos.empty()) {
      return;
    }
    if constexpr (Architecture == lbmpy::Arch::CPU) {
      for (std::size_t i = 0ul; i < pos.size(); ++i) {
        add_force_at_pos(pos[i], forces[i]);
      }
    }
#if defined(__CUDACC__)
    if constexpr (Architecture == lbmpy::Arch::GPU) {
      auto const &lattice = get_lattice();
      auto const &block = *(lattice.get_blocks()->begin());
      auto const origin = block.getAABB().min();
      std::vector<FloatType> host_pos;
      std::vector<FloatType> host_force;
      host_pos.reserve(3ul * pos.size());
      host_force.reserve(3ul * forces.size());
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      for (auto const &vec : pos) {
#pragma unroll
        for (std::size_t i : {0ul, 1ul, 2ul}) {
          host_pos.emplace_back(static_cast<FloatType>(vec[i] - origin[i]));
        }
      }
      for (auto const &vec : forces) {
#pragma unroll
        for (std::size_t i : {0ul, 1ul, 2ul}) {
          host_force.emplace_back(static_cast<FloatType>(vec[i]));
        }
      }
      auto const gl = lattice.get_ghost_layers();
      auto field = block.template uncheckedFastGetData<VectorField>(
          m_force_to_be_applied_id);
      lbm::accessor::Interpolation::set(field, host_pos, host_force, gl);
    }
#endif
  }
 
  std::vector<Utils::Vector3d>
  get_velocities_at_pos(std::vector<Utils::Vector3d> const &pos) override {
    if (pos.empty()) {
      return {};
    }
    if constexpr (Architecture == lbmpy::Arch::CPU) {
      std::vector<Utils::Vector3d> vel{};
      vel.reserve(pos.size());
      for (auto const &vec : pos) {
        auto res = get_velocity_at_pos(vec, true);
        assert(res.has_value());
        vel.emplace_back(*res);
      }
      return vel;
    }
#if defined(__CUDACC__)
    if constexpr (Architecture == lbmpy::Arch::GPU) {
      auto const &lattice = get_lattice();
      auto const &block = *(lattice.get_blocks()->begin());
      auto const origin = block.getAABB().min();
      std::vector<FloatType> host_pos;
      host_pos.reserve(3ul * pos.size());
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      for (auto const &vec : pos) {
#pragma unroll
        for (std::size_t i : {0ul, 1ul, 2ul}) {
          host_pos.emplace_back(static_cast<FloatType>(vec[i] - origin[i]));
        }
      }
      auto const gl = lattice.get_ghost_layers();
      auto field =
          block.template uncheckedFastGetData<VectorField>(m_velocity_field_id);
      auto const res = lbm::accessor::Interpolation::get(field, host_pos, gl);
      std::vector<Utils::Vector3d> vel{};
      vel.reserve(res.size() / 3ul);
      for (auto it = res.begin(); it != res.end(); it += 3) {
        vel.emplace_back(Utils::Vector3d{static_cast<double>(*(it + 0)),
                                         static_cast<double>(*(it + 1)),
                                         static_cast<double>(*(it + 2))});
      }
      return vel;
    }
#endif
    return {};
  }
 
  std::optional<Utils::Vector3d>
  get_velocity_at_pos(Utils::Vector3d const &pos,
                      bool consider_points_in_halo = false) const override {
    if (!consider_points_in_halo and !m_lattice->pos_in_local_domain(pos))
      return std::nullopt;
    if (consider_points_in_halo and !m_lattice->pos_in_local_halo(pos))
      return std::nullopt;
    Utils::Vector3d v{0., 0., 0.};
    interpolate_bspline_at_pos(
        pos, [this, &v, &pos](std::array<int, 3> const node, double weight) {
          // Nodes with zero weight might not be accessible, because they can be
          // outside ghost layers
          if (weight != 0.) {
            auto const res = get_node_velocity(Utils::Vector3i(node), true);
            if (!res) {
              throw interpolation_illegal_access("velocity", pos, node, weight);
            }
            v += *res * weight;
          }
        });
    return {std::move(v)};
  }
 
  std::optional<double>
  get_density_at_pos(Utils::Vector3d const &pos,
                     bool consider_points_in_halo = false) const override {
    if (!consider_points_in_halo and !m_lattice->pos_in_local_domain(pos))
      return std::nullopt;
    if (consider_points_in_halo and !m_lattice->pos_in_local_halo(pos))
      return std::nullopt;
    double dens = 0.;
    interpolate_bspline_at_pos(
        pos, [this, &dens, &pos](std::array<int, 3> const node, double weight) {
          // Nodes with zero weight might not be accessible, because they can be
          // outside ghost layers
          if (weight != 0.) {
            auto const res = get_node_density(Utils::Vector3i(node), true);
            if (!res) {
              throw interpolation_illegal_access("density", pos, node, weight);
            }
            dens += *res * weight;
          }
        });
    return {std::move(dens)};
  }
 
  // Local force
  bool add_force_at_pos(Utils::Vector3d const &pos,
                        Utils::Vector3d const &force) override {
    if (!m_lattice->pos_in_local_halo(pos))
      return false;
    auto const force_at_node = [this, &force](std::array<int, 3> const node,
                                              double weight) {
      auto const bc =
          get_block_and_cell(get_lattice(), Utils::Vector3i(node), true);
      if (bc) {
        auto const weighted_force = to_vector3<FloatType>(weight * force);
        auto force_field =
            bc->block->template uncheckedFastGetData<VectorField>(
                m_force_to_be_applied_id);
        lbm::accessor::Vector::add(force_field, weighted_force, bc->cell);
      }
    };
    interpolate_bspline_at_pos(pos, force_at_node);
    return true;
  }
 
  std::optional<Utils::Vector3d>
  get_node_force_to_be_applied(Utils::Vector3i const &node) const override {
    auto const bc = get_block_and_cell(get_lattice(), node, true);
    if (!bc)
      return std::nullopt;
 
    auto field =
        bc->block->template getData<VectorField>(m_force_to_be_applied_id);
    auto const vec = lbm::accessor::Vector::get(field, bc->cell);
    return to_vector3d(vec);
  }
 
  std::optional<Utils::Vector3d>
  get_node_last_applied_force(Utils::Vector3i const &node,
                              bool consider_ghosts = false) const override {
    auto const bc = get_block_and_cell(get_lattice(), node, consider_ghosts);
    if (!bc)
      return std::nullopt;
 
    auto const field =
        bc->block->template getData<VectorField>(m_last_applied_force_field_id);
    auto const vec = lbm::accessor::Vector::get(field, bc->cell);
    return to_vector3d(vec);
  }
 
  bool set_node_last_applied_force(Utils::Vector3i const &node,
                                   Utils::Vector3d const &force) override {
    auto bc = get_block_and_cell(get_lattice(), node, false);
    if (!bc)
      return false;
 
    auto field =
        bc->block->template getData<VectorField>(m_last_applied_force_field_id);
    auto const vec = to_vector3<FloatType>(force);
    lbm::accessor::Vector::set(field, vec, bc->cell);
 
    return true;
  }
 
  std::vector<double> get_slice_last_applied_force(
      Utils::Vector3i const &lower_corner,
      Utils::Vector3i const &upper_corner) const override {
    std::vector<double> out;
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto const &block = *(lattice.get_blocks()->begin());
      auto const field =
          block.template getData<VectorField>(m_last_applied_force_field_id);
      auto const lower_cell = ci->min();
      auto const upper_cell = ci->max();
      auto const n_values = 3u * ci->numCells();
      out.reserve(n_values);
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      for (auto x = lower_cell.x(); x <= upper_cell.x(); ++x) {
        for (auto y = lower_cell.y(); y <= upper_cell.y(); ++y) {
          for (auto z = lower_cell.z(); z <= upper_cell.z(); ++z) {
            auto const vec = lbm::accessor::Vector::get(field, Cell{x, y, z});
            for (uint_t f = 0u; f < 3u; ++f) {
              out.emplace_back(double_c(vec[f]));
            }
          }
        }
      }
      assert(out.size() == n_values);
    }
    return out;
  }
 
  void set_slice_last_applied_force(Utils::Vector3i const &lower_corner,
                                    Utils::Vector3i const &upper_corner,
                                    std::vector<double> const &force) override {
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto &block = *(lattice.get_blocks()->begin());
      auto field =
          block.template getData<VectorField>(m_last_applied_force_field_id);
      auto const lower_cell = ci->min();
      auto const upper_cell = ci->max();
      auto it = force.begin();
      assert(force.size() == 3u * ci->numCells());
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      for (auto x = lower_cell.x(); x <= upper_cell.x(); ++x) {
        for (auto y = lower_cell.y(); y <= upper_cell.y(); ++y) {
          for (auto z = lower_cell.z(); z <= upper_cell.z(); ++z) {
            Vector3<FloatType> vec;
            for (uint_t f = 0u; f < 3u; ++f) {
              vec[f] = FloatType_c(*it);
              ++it;
            }
            lbm::accessor::Vector::set(field, vec, Cell{x, y, z});
          }
        }
      }
    }
  }
 
  // Population
  std::optional<std::vector<double>>
  get_node_population(Utils::Vector3i const &node,
                      bool consider_ghosts = false) const override {
    auto bc = get_block_and_cell(get_lattice(), node, consider_ghosts);
    if (!bc)
      return std::nullopt;
 
    auto pdf_field = bc->block->template getData<PdfField>(m_pdf_field_id);
    auto const pop = lbm::accessor::Population::get(pdf_field, bc->cell);
    std::vector<double> population(Stencil::Size);
    for (uint_t f = 0u; f < Stencil::Size; ++f) {
      population[f] = double_c(pop[f]);
    }
 
    return {std::move(population)};
  }
 
  bool set_node_population(Utils::Vector3i const &node,
                           std::vector<double> const &population) override {
    auto bc = get_block_and_cell(get_lattice(), node, false);
    if (!bc)
      return false;
 
    auto pdf_field = bc->block->template getData<PdfField>(m_pdf_field_id);
    std::array<FloatType, Stencil::Size> pop;
    for (uint_t f = 0u; f < Stencil::Size; ++f) {
      pop[f] = FloatType_c(population[f]);
    }
    lbm::accessor::Population::set(pdf_field, pop, bc->cell);
 
    return true;
  }
 
  std::vector<double>
  get_slice_population(Utils::Vector3i const &lower_corner,
                       Utils::Vector3i const &upper_corner) const override {
    std::vector<double> out;
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto const &block = *(lattice.get_blocks()->begin());
      auto const pdf_field = block.template getData<PdfField>(m_pdf_field_id);
      auto const values = lbm::accessor::Population::get(pdf_field, *ci);
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      if constexpr (std::is_same_v<typename decltype(values)::value_type,
                                   double>) {
        out = std::move(values);
      } else {
        out = std::vector<double>(values.begin(), values.end());
      }
      assert(out.size() == stencil_size() * ci->numCells());
    }
    return out;
  }
 
  void set_slice_population(Utils::Vector3i const &lower_corner,
                            Utils::Vector3i const &upper_corner,
                            std::vector<double> const &population) override {
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto &block = *(lattice.get_blocks()->begin());
      auto pdf_field = block.template getData<PdfField>(m_pdf_field_id);
      assert(population.size() == stencil_size() * ci->numCells());
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      std::vector<FloatType> const values(population.begin(), population.end());
      lbm::accessor::Population::set(pdf_field, values, *ci);
    }
  }
 
  // Density
  std::optional<double>
  get_node_density(Utils::Vector3i const &node,
                   bool consider_ghosts = false) const override {
    auto bc = get_block_and_cell(get_lattice(), node, consider_ghosts);
    if (!bc)
      return std::nullopt;
 
    auto pdf_field =
        bc->block->template uncheckedFastGetData<PdfField>(m_pdf_field_id);
    auto const density = lbm::accessor::Density::get(pdf_field, bc->cell);
    return {double_c(density)};
  }
 
  bool set_node_density(Utils::Vector3i const &node, double density) override {
    auto bc = get_block_and_cell(get_lattice(), node, false);
    if (!bc)
      return false;
 
    auto pdf_field = bc->block->template getData<PdfField>(m_pdf_field_id);
    lbm::accessor::Density::set(pdf_field, FloatType_c(density), bc->cell);
 
    return true;
  }
 
  std::vector<double>
  get_slice_density(Utils::Vector3i const &lower_corner,
                    Utils::Vector3i const &upper_corner) const override {
    std::vector<double> out;
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto const &block = *(lattice.get_blocks()->begin());
      auto const pdf_field = block.template getData<PdfField>(m_pdf_field_id);
      auto const values = lbm::accessor::Density::get(pdf_field, *ci);
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      if constexpr (std::is_same_v<typename decltype(values)::value_type,
                                   double>) {
        out = std::move(values);
      } else {
        out = std::vector<double>(values.begin(), values.end());
      }
      assert(out.size() == ci->numCells());
    }
    return out;
  }
 
  void set_slice_density(Utils::Vector3i const &lower_corner,
                         Utils::Vector3i const &upper_corner,
                         std::vector<double> const &density) override {
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto &block = *(lattice.get_blocks()->begin());
      auto pdf_field = block.template getData<PdfField>(m_pdf_field_id);
      assert(density.size() == ci->numCells());
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      std::vector<FloatType> const values(density.begin(), density.end());
      lbm::accessor::Density::set(pdf_field, values, *ci);
    }
  }
 
  std::optional<Utils::Vector3d>
  get_node_velocity_at_boundary(Utils::Vector3i const &node,
                                bool consider_ghosts = false) const override {
    auto const bc = get_block_and_cell(get_lattice(), node, consider_ghosts);
    if (!bc or !m_boundary->node_is_boundary(node))
      return std::nullopt;
 
    return {to_vector3d(m_boundary->get_node_value_at_boundary(node))};
  }
 
  bool set_node_velocity_at_boundary(Utils::Vector3i const &node,
                                     Utils::Vector3d const &velocity) override {
    auto bc = get_block_and_cell(get_lattice(), node, true);
    if (bc) {
      m_boundary->set_node_value_at_boundary(
          node, to_vector3<FloatType>(velocity), *bc);
    }
    return bc.has_value();
  }
 
  std::vector<std::optional<Utils::Vector3d>> get_slice_velocity_at_boundary(
      Utils::Vector3i const &lower_corner,
      Utils::Vector3i const &upper_corner) const override {
    std::vector<std::optional<Utils::Vector3d>> out;
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto const local_offset = std::get<0>(lattice.get_local_grid_range());
      auto const lower_cell = ci->min();
      auto const upper_cell = ci->max();
      auto const n_values = ci->numCells();
      out.reserve(n_values);
      for (auto x = lower_cell.x(); x <= upper_cell.x(); ++x) {
        for (auto y = lower_cell.y(); y <= upper_cell.y(); ++y) {
          for (auto z = lower_cell.z(); z <= upper_cell.z(); ++z) {
            auto const node = local_offset + Utils::Vector3i{{x, y, z}};
            if (m_boundary->node_is_boundary(node)) {
              out.emplace_back(
                  to_vector3d(m_boundary->get_node_value_at_boundary(node)));
            } else {
              out.emplace_back(std::nullopt);
            }
          }
        }
      }
      assert(out.size() == n_values);
    }
    return out;
  }
 
  void set_slice_velocity_at_boundary(
      Utils::Vector3i const &lower_corner, Utils::Vector3i const &upper_corner,
      std::vector<std::optional<Utils::Vector3d>> const &velocity) override {
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto const local_offset = std::get<0>(lattice.get_local_grid_range());
      auto const lower_cell = ci->min();
      auto const upper_cell = ci->max();
      auto it = velocity.begin();
      assert(velocity.size() == ci->numCells());
      for (auto x = lower_cell.x(); x <= upper_cell.x(); ++x) {
        for (auto y = lower_cell.y(); y <= upper_cell.y(); ++y) {
          for (auto z = lower_cell.z(); z <= upper_cell.z(); ++z) {
            auto const node = local_offset + Utils::Vector3i{{x, y, z}};
            auto const bc = get_block_and_cell(lattice, node, false);
            auto const &opt = *it;
            if (opt) {
              m_boundary->set_node_value_at_boundary(
                  node, to_vector3<FloatType>(*opt), *bc);
            } else {
              m_boundary->remove_node_from_boundary(node, *bc);
            }
            ++it;
          }
        }
      }
    }
  }
 
  std::optional<Utils::Vector3d>
  get_node_boundary_force(Utils::Vector3i const &node) const override {
    auto const bc = get_block_and_cell(get_lattice(), node, true);
    if (!bc or !m_boundary->node_is_boundary(node))
      return std::nullopt;
 
    return get_node_last_applied_force(node, true);
  }
 
  bool remove_node_from_boundary(Utils::Vector3i const &node) override {
    auto bc = get_block_and_cell(get_lattice(), node, true);
    if (bc) {
      m_boundary->remove_node_from_boundary(node, *bc);
    }
    return bc.has_value();
  }
 
  std::optional<bool>
  get_node_is_boundary(Utils::Vector3i const &node,
                       bool consider_ghosts = false) const override {
    auto const bc = get_block_and_cell(get_lattice(), node, consider_ghosts);
    if (!bc)
      return std::nullopt;
 
    return {m_boundary->node_is_boundary(node)};
  }
 
  std::vector<bool>
  get_slice_is_boundary(Utils::Vector3i const &lower_corner,
                        Utils::Vector3i const &upper_corner) const override {
    std::vector<bool> out;
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto const local_offset = std::get<0>(lattice.get_local_grid_range());
      auto const lower_cell = ci->min();
      auto const upper_cell = ci->max();
      auto const n_values = ci->numCells();
      out.reserve(n_values);
      for (auto x = lower_cell.x(); x <= upper_cell.x(); ++x) {
        for (auto y = lower_cell.y(); y <= upper_cell.y(); ++y) {
          for (auto z = lower_cell.z(); z <= upper_cell.z(); ++z) {
            auto const node = local_offset + Utils::Vector3i{x, y, z};
            out.emplace_back(m_boundary->node_is_boundary(node));
          }
        }
      }
      assert(out.size() == n_values);
    }
    return out;
  }
 
  void reallocate_ubb_field() override { m_boundary->boundary_update(); }
 
  void clear_boundaries() override {
    reset_boundary_handling();
    ghost_communication();
  }
 
  void
  update_boundary_from_shape(std::vector<int> const &raster_flat,
                             std::vector<double> const &data_flat) override {
    auto const grid_size = get_lattice().get_grid_dimensions();
    auto const data = fill_3D_vector_array(data_flat, grid_size);
    set_boundary_from_grid(*m_boundary, get_lattice(), raster_flat, data);
    ghost_communication();
    reallocate_ubb_field();
  }
 
  // Pressure tensor
  std::optional<Utils::VectorXd<9>>
  get_node_pressure_tensor(Utils::Vector3i const &node) const override {
    auto bc = get_block_and_cell(get_lattice(), node, false);
    if (!bc)
      return std::nullopt;
 
    auto pdf_field = bc->block->template getData<PdfField>(m_pdf_field_id);
    auto tensor = lbm::accessor::PressureTensor::get(pdf_field, bc->cell);
    pressure_tensor_correction(tensor);
 
    return to_vector9d(tensor);
  }
 
  std::vector<double> get_slice_pressure_tensor(
      Utils::Vector3i const &lower_corner,
      Utils::Vector3i const &upper_corner) const override {
    std::vector<double> out;
    if (auto const ci = get_interval(lower_corner, upper_corner)) {
      auto const &lattice = get_lattice();
      auto const &block = *(lattice.get_blocks()->begin());
      auto const pdf_field = block.template getData<PdfField>(m_pdf_field_id);
      auto const lower_cell = ci->min();
      auto const upper_cell = ci->max();
      auto const n_values = 9u * ci->numCells();
      out.reserve(n_values);
      assert(++(lattice.get_blocks()->begin()) == lattice.get_blocks()->end());
      for (auto x = lower_cell.x(); x <= upper_cell.x(); ++x) {
        for (auto y = lower_cell.y(); y <= upper_cell.y(); ++y) {
          for (auto z = lower_cell.z(); z <= upper_cell.z(); ++z) {
            auto const cell = Cell{x, y, z};
            auto tensor = lbm::accessor::PressureTensor::get(pdf_field, cell);
            pressure_tensor_correction(tensor);
            for (auto i = 0u; i < 9u; ++i) {
              out.emplace_back(tensor[i]);
            }
          }
        }
      }
      assert(out.size() == n_values);
    }
    return out;
  }
 
  // Global pressure tensor
  [[nodiscard]] Utils::VectorXd<9> get_pressure_tensor() const override {
    auto const &blocks = get_lattice().get_blocks();
    Matrix3<FloatType> tensor(FloatType{0});
    for (auto block = blocks->begin(); block != blocks->end(); ++block) {
      auto pdf_field = block->template getData<PdfField>(m_pdf_field_id);
      WALBERLA_FOR_ALL_CELLS_XYZ(pdf_field, {
        tensor += lbm::accessor::PressureTensor::get(pdf_field, Cell{x, y, z});
      });
    }
    auto const grid_size = get_lattice().get_grid_dimensions();
    auto const number_of_nodes = Utils::product(grid_size);
    pressure_tensor_correction(tensor);
    return to_vector9d(tensor) * (1. / static_cast<double>(number_of_nodes));
  }
 
  // Global momentum
  [[nodiscard]] Utils::Vector3d get_momentum() const override {
    auto const &blocks = get_lattice().get_blocks();
    Vector3<FloatType> mom(FloatType{0});
    for (auto block = blocks->begin(); block != blocks->end(); ++block) {
      auto pdf_field = block->template getData<PdfField>(m_pdf_field_id);
      auto force_field =
          block->template getData<VectorField>(m_last_applied_force_field_id);
      mom += lbm::accessor::MomentumDensity::reduce(pdf_field, force_field);
    }
    return to_vector3d(mom);
  }
 
  // Global external force
  void set_external_force(Utils::Vector3d const &ext_force) override {
    m_reset_force->set_ext_force(ext_force);
  }
 
  [[nodiscard]] Utils::Vector3d get_external_force() const noexcept override {
    return m_reset_force->get_ext_force();
  }
 
  [[nodiscard]] double get_kT() const noexcept override {
    return numeric_cast<double>(m_kT);
  }
 
  [[nodiscard]] std::optional<uint64_t> get_rng_state() const override {
    auto const cm = std::get_if<CollisionModelThermalized>(&*m_collision_model);
    if (!cm or m_kT == 0.) {
      return std::nullopt;
    }
    return {static_cast<uint64_t>(cm->time_step_)};
  }
 
  void set_rng_state(uint64_t counter) override {
    auto const cm = std::get_if<CollisionModelThermalized>(&*m_collision_model);
    if (!cm or m_kT == 0.) {
      throw std::runtime_error("This LB instance is unthermalized");
    }
    assert(counter <=
           static_cast<uint32_t>(std::numeric_limits<uint_t>::max()));
    cm->time_step_ = static_cast<uint32_t>(counter);
  }
 
  [[nodiscard]] LatticeWalberla const &get_lattice() const noexcept override {
    return *m_lattice;
  }
 
  [[nodiscard]] std::size_t get_velocity_field_id() const noexcept override {
    return m_velocity_field_id;
  }
 
  [[nodiscard]] std::size_t get_force_field_id() const noexcept override {
    return m_force_to_be_applied_id;
  }
 
  void register_vtk_field_filters(walberla::vtk::VTKOutput &vtk_obj) override {
    if constexpr (Architecture == lbmpy::Arch::GPU) {
      throw std::runtime_error("VTK output not supported for GPU");
    } else {
      field::FlagFieldCellFilter<FlagField> fluid_filter(m_flag_field_id);
      fluid_filter.addFlag(Boundary_flag);
      vtk_obj.addCellExclusionFilter(fluid_filter);
    }
  }
 
protected:
  template <typename Field_T, uint_t F_SIZE_ARG, typename OutputType>
  class VTKWriter : public vtk::BlockCellDataWriter<OutputType, F_SIZE_ARG> {
  public:
    VTKWriter(ConstBlockDataID const &block_id, std::string const &id,
              FloatType unit_conversion)
        : vtk::BlockCellDataWriter<OutputType, F_SIZE_ARG>(id),
          m_block_id(block_id), m_field(nullptr),
          m_conversion(unit_conversion) {}
 
  protected:
    void configure() override {
      WALBERLA_ASSERT_NOT_NULLPTR(this->block_);
      m_field = this->block_->template getData<Field_T>(m_block_id);
    }
 
    ConstBlockDataID const m_block_id;
    Field_T const *m_field;
    FloatType const m_conversion;
  };
 
  template <typename OutputType = float>
  class DensityVTKWriter : public VTKWriter<PdfField, 1u, OutputType> {
  public:
    using VTKWriter<PdfField, 1u, OutputType>::VTKWriter;
 
  protected:
    OutputType evaluate(cell_idx_t const x, cell_idx_t const y,
                        cell_idx_t const z, cell_idx_t const) override {
      WALBERLA_ASSERT_NOT_NULLPTR(this->m_field);
      auto const density =
          lbm::accessor::Density::get(this->m_field, {x, y, z});
      return numeric_cast<OutputType>(this->m_conversion * density);
    }
  };
 
  template <typename OutputType = float>
  class VelocityVTKWriter : public VTKWriter<VectorField, 3u, OutputType> {
  public:
    using VTKWriter<VectorField, 3u, OutputType>::VTKWriter;
 
  protected:
    OutputType evaluate(cell_idx_t const x, cell_idx_t const y,
                        cell_idx_t const z, cell_idx_t const f) override {
      WALBERLA_ASSERT_NOT_NULLPTR(this->m_field);
      auto const velocity =
          lbm::accessor::Vector::get(this->m_field, {x, y, z});
      return numeric_cast<OutputType>(this->m_conversion * velocity[uint_c(f)]);
    }
  };
 
  template <typename OutputType = float>
  class PressureTensorVTKWriter : public VTKWriter<PdfField, 9u, OutputType> {
  public:
    PressureTensorVTKWriter(ConstBlockDataID const &block_id,
                            std::string const &id, FloatType unit_conversion,
                            FloatType off_diag_factor)
        : VTKWriter<PdfField, 9u, OutputType>::VTKWriter(block_id, id,
                                                         unit_conversion),
          m_off_diag_factor(off_diag_factor) {}
 
  protected:
    OutputType evaluate(cell_idx_t const x, cell_idx_t const y,
                        cell_idx_t const z, cell_idx_t const f) override {
      WALBERLA_ASSERT_NOT_NULLPTR(this->m_field);
      auto const pressure =
          lbm::accessor::PressureTensor::get(this->m_field, {x, y, z});
      auto const revert_factor =
          (f == 0 or f == 4 or f == 8) ? FloatType{1} : m_off_diag_factor;
      return numeric_cast<OutputType>(this->m_conversion * revert_factor *
                                      pressure[uint_c(f)]);
    }
    FloatType const m_off_diag_factor;
  };
 
public:
  void register_vtk_field_writers(walberla::vtk::VTKOutput &vtk_obj,
                                  LatticeModel::units_map const &units,
                                  int flag_observables) override {
    if constexpr (Architecture == lbmpy::Arch::GPU) {
      throw std::runtime_error("VTK output not supported for GPU");
    } else {
      if (flag_observables & static_cast<int>(OutputVTK::density)) {
        auto const unit_conversion = FloatType_c(units.at("density"));
        vtk_obj.addCellDataWriter(make_shared<DensityVTKWriter<float>>(
            m_pdf_field_id, "density", unit_conversion));
      }
      if (flag_observables & static_cast<int>(OutputVTK::velocity_vector)) {
        auto const unit_conversion = FloatType_c(units.at("velocity"));
        vtk_obj.addCellDataWriter(make_shared<VelocityVTKWriter<float>>(
            m_velocity_field_id, "velocity_vector", unit_conversion));
      }
      if (flag_observables & static_cast<int>(OutputVTK::pressure_tensor)) {
        auto const unit_conversion = FloatType_c(units.at("pressure"));
        vtk_obj.addCellDataWriter(make_shared<PressureTensorVTKWriter<float>>(
            m_pdf_field_id, "pressure_tensor", unit_conversion,
            pressure_tensor_correction_factor()));
      }
    }
  }
 
  ~LBWalberlaImpl() override = default;
};
 
} // namespace walberla