LBWalberlaImpl.hpp

Go to the documentation of this file.

/*
 * Copyright (C) 2019-2026 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 <domain_decomposition/BlockDataID.h>
#include <domain_decomposition/IBlock.h>
#include <field/AddToStorage.h>
#include <field/vtk/FlagFieldCellFilter.h>
#include <field/vtk/VTKWriter.h>
#include <stencil/D3Q19.h>
#include <stencil/D3Q27.h>
#include <waLBerlaDefinitions.h>
#if defined(__CUDACC__) and defined(WALBERLA_BUILD_WITH_CUDA)
#include <gpu/AddGPUFieldToStorage.h>
#include <gpu/HostFieldAllocator.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_fields.hpp"
#include "lb_kernels.hpp"
#if defined(__CUDACC__) and defined(WALBERLA_BUILD_WITH_CUDA)
#include "lb_fields.cuh"
#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 <walberla_bridge/utils/ResourceManager.hpp>
#include <walberla_bridge/walberla_init.hpp>
 
#include <utils/Vector.hpp>
 
#include <array>
#include <bitset>
#include <cstddef>
#include <functional>
#include <initializer_list>
#include <limits>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#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 {
#if not defined(WALBERLA_BUILD_WITH_CUDA)
  static_assert(Architecture != lbmpy::Arch::GPU,
                "waLBerla was compiled without CUDA support");
#endif
protected:
  // ---- Types & Constants ----
 
  using Kernels = detail::KernelTrait<FloatType, Architecture>;
  using BoundaryModel = BoundaryHandling<FloatType, Vector3<FloatType>,
                                         typename Kernels::DynamicUBB>;
  using CollisionModel =
      std::variant<typename Kernels::StreamCollisionModelThermalized,
                   typename Kernels::StreamCollisionModelLeesEdwards>;
 
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 BlockStorage = LatticeWalberla::Lattice_T;
 
protected:
  // "underlying" field types (`GPUField` has no f-size info at compile time)
  using _PdfField = FieldTrait<FloatType, Stencil>::PdfField;
  using _VectorField = FieldTrait<FloatType, Stencil>::VectorField;
 
public:
  using PdfField = FieldTrait<FloatType, Stencil, Architecture>::PdfField;
  using VectorField = FieldTrait<FloatType, Stencil, Architecture>::VectorField;
  using FlagField = BoundaryModel::FlagField;
#if defined(__CUDACC__) and defined(WALBERLA_BUILD_WITH_CUDA)
  using GPUField = gpu::GPUField<FloatType>;
  using PdfFieldCpu =
      FieldTrait<FloatType, Stencil, lbmpy::Arch::CPU>::PdfField;
  using VectorFieldCpu =
      FieldTrait<FloatType, Stencil, lbmpy::Arch::CPU>::VectorField;
#endif
 
  struct GhostComm {
    /** @brief Ghost communication operations. */
    enum GhostCommFlags : unsigned {
      PDF, ///< PDFs communication
      VEL, ///< velocities communication
      LAF, ///< last applied forces communication
      UBB, ///< boundaries communication
      SIZE
    };
  };
 
protected:
  /**
   * @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 =
      FieldTrait<FloatType, Stencil,
                 Architecture>::template RegularCommScheme<stencil::D3Q27>;
  using BoundaryFullCommunicator =
      FieldTrait<FloatType, Stencil,
                 Architecture>::template BoundaryCommScheme<stencil::D3Q27>;
  /**
   * @brief Regular communicator.
   * We use the same directions as the stencil during integration.
   */
  using PDFStreamingCommunicator =
      FieldTrait<FloatType, Stencil,
                 Architecture>::template RegularCommScheme<Stencil>;
  template <class Field>
  using PackInfo =
      FieldTrait<FloatType, Stencil, Architecture>::template PackInfo<Field>;
 
protected:
  // ---- Member Variables ----
 
  // Physical parameters
  FloatType m_viscosity; /// kinematic viscosity
  FloatType m_density;
  FloatType m_kT;
  unsigned int m_seed;
  double m_zc_to_md; // zero-centered conversion factor to MD units
  double m_zc_to_lb; // zero-centered conversion factor to LB units
 
  // lattice
  std::shared_ptr<LatticeWalberla> m_lattice;
 
  // 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_vel_tmp_field_id;
 
#if defined(__CUDACC__) and defined(WALBERLA_BUILD_WITH_CUDA)
  std::optional<BlockDataID> m_pdf_cpu_field_id;
  std::optional<BlockDataID> m_vel_cpu_field_id;
#endif
 
  /** Flag for boundary cells. */
  FlagUID const Boundary_flag{"boundary"};
  bool m_has_boundaries{false};
 
  // boundaries
  std::shared_ptr<BoundaryModel> m_boundary;
 
  // communicators
  std::shared_ptr<BoundaryFullCommunicator> m_boundary_communicator;
  std::shared_ptr<RegularFullCommunicator> m_full_communicator;
  std::shared_ptr<RegularFullCommunicator> m_pdf_communicator;
  std::shared_ptr<RegularFullCommunicator> m_vel_communicator;
  std::shared_ptr<RegularFullCommunicator> m_laf_communicator;
  std::shared_ptr<PDFStreamingCommunicator> m_pdf_streaming_communicator;
  std::bitset<GhostComm::SIZE> m_pending_ghost_comm;
  ResourceObserver m_mpi_cart_comm_observer;
 
  // collision sweep
  std::shared_ptr<CollisionModel> m_collision_model;
 
  // force reset sweep + external force handling
  std::shared_ptr<ResetForce<PdfField, VectorField>> m_reset_force;
 
  // velocity update sweep
  std::shared_ptr<typename Kernels::UpdateVelFromPDF>
      m_update_velocities_from_pdf;
 
  // 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;
 
#if defined(__CUDACC__) and defined(WALBERLA_BUILD_WITH_CUDA)
  std::shared_ptr<gpu::HostFieldAllocator<FloatType>> m_host_field_allocator;
#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]] bool is_double_precision() const noexcept override {
    return std::is_same_v<FloatType, double>;
  }
 
  [[nodiscard]] bool is_gpu() const noexcept override {
    return Architecture == lbmpy::Arch::GPU;
  }
 
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_seed(0u), m_zc_to_md(density),
        m_zc_to_lb(1. / density), m_lattice(std::move(lattice)),
        m_mpi_cart_comm_observer(get_mpi_cart_comm_observer()) {
 
    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_vel_tmp_field_id = add_to_storage<_VectorField>("velocity_tmp");
#if defined(__CUDACC__) and defined(WALBERLA_BUILD_WITH_CUDA)
    m_host_field_allocator =
        std::make_shared<gpu::HostFieldAllocator<FloatType>>();
#endif
 
    // Initialize and register pdf field with zero centered density
    auto pdf_setter = typename Kernels::InitialPDFsSetter(
        m_force_to_be_applied_id, m_pdf_field_id, m_velocity_field_id, 1.0);
    for (auto &block : *blocks) {
      pdf_setter(&block);
    }
 
    // 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(m_lattice->get_blocks());
 
    // Set up the communication and register fields
    setup_streaming_communicator();
    setup_full_communicator();
 
    m_pending_ghost_comm.set();
 
    // 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);
 
    // Instantiate velocity update sweep
    m_update_velocities_from_pdf =
        std::make_shared<typename Kernels::UpdateVelFromPDF>(
            m_last_applied_force_field_id, m_pdf_field_id, m_velocity_field_id);
  }
 
  ~LBWalberlaImpl() override = default;
 
  // ---- Integration (Core LB Algorithm) ----
 
  void integrate() override {
    integrate_pull_scheme();
    integrate_vtk_writers();
  }
 
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++;
      }
    }
  }
 
private:
  /**
   * @brief One LB time step using the pull scheme.
   * Sequence: reset forces, stream-collide, communicate PDFs,
   * apply Lees-Edwards interpolation (if active), handle boundaries,
   * update velocity field from PDFs.
   */
  void integrate_pull_scheme() {
    assert(m_mpi_cart_comm_observer.is_valid());
    auto const &blocks = get_lattice().get_blocks();
    // Reset force fields
    integrate_reset_force(blocks);
    // LB stream collide
    integrate_stream_collide(blocks);
    // Mark pending ghost layer updates
    // As pdf and laf are communicated directly afterwards, they are not set
    m_pending_ghost_comm.set(GhostComm::VEL);
    m_pending_ghost_comm.set(GhostComm::LAF);
    m_pdf_streaming_communicator->communicate();
    if (has_lees_edwards_bc()) {
      apply_lees_edwards_pdf_interpolation(blocks);
      apply_lees_edwards_last_applied_force_interpolation(blocks);
    }
    // Handle boundaries
    if (m_has_boundaries) {
      integrate_boundaries(blocks);
    }
    // Update velocities from pdfs
    integrate_update_velocities_from_pdf(blocks);
 
    if (has_lees_edwards_bc()) {
      apply_lees_edwards_vel_interpolation_and_shift(blocks);
    }
  }
 
  void integrate_stream_collide(std::shared_ptr<BlockStorage> const &blocks) {
    auto &cm_variant = *m_collision_model;
    for (auto &block : *blocks) {
      auto const block_variant = std::variant<IBlock *>(&block);
      std::visit(m_run_stream_collide_sweep, cm_variant, block_variant);
    }
    if (auto *cm =
            std::get_if<typename Kernels::StreamCollisionModelThermalized>(
                &cm_variant)) {
      cm->setTime_step(cm->getTime_step() + 1u);
    }
  }
 
  void integrate_reset_force(std::shared_ptr<BlockStorage> const &blocks) {
    for (auto &block : *blocks)
      (*m_reset_force)(&block);
  }
 
  void integrate_boundaries(std::shared_ptr<BlockStorage> const &blocks) {
    for (auto &block : *blocks)
      (*m_boundary)(&block);
  }
 
  void integrate_update_velocities_from_pdf(
      std::shared_ptr<BlockStorage> const &blocks) {
    for (auto &block : *blocks)
      (*m_update_velocities_from_pdf)(&block);
  }
 
private:
  // ---- Collision Model ----
 
  /**
   * @brief Visitor for dispatching stream-collide sweeps.
   * Handles both thermalized and Lees-Edwards collision models
   * via @c std::visit on the @ref CollisionModel variant.
   */
  class StreamCollideSweepVisitor {
  public:
    using StructuredBlockStorage = LatticeWalberla::Lattice_T;
 
    void operator()(typename Kernels::StreamCollisionModelThermalized &cm,
                    IBlock *b) {
      cm.configure(m_storage, b);
      cm(b);
    }
 
    void operator()(typename Kernels::StreamCollisionModelLeesEdwards &cm,
                    IBlock *b) {
      cm.setV_s(static_cast<decltype(cm.getV_s())>(
          m_lees_edwards_callbacks->get_shear_velocity()));
      cm(b);
    }
 
    StreamCollideSweepVisitor() = default;
    StreamCollideSweepVisitor(std::shared_ptr<StructuredBlockStorage> storage) {
      m_storage = std::move(storage);
    }
    StreamCollideSweepVisitor(std::shared_ptr<StructuredBlockStorage> storage,
                              std::shared_ptr<LeesEdwardsPack> callbacks) {
      m_storage = std::move(storage);
      m_lees_edwards_callbacks = std::move(callbacks);
    }
 
  private:
    std::shared_ptr<StructuredBlockStorage> m_storage{};
    std::shared_ptr<LeesEdwardsPack> m_lees_edwards_callbacks{};
  };
  StreamCollideSweepVisitor m_run_stream_collide_sweep{};
 
  /** @brief Relaxation rate omega from kinematic viscosity: 2/(6*nu+1). */
  FloatType shear_mode_relaxation_rate() const;
  /**
   * @brief Odd-mode relaxation rate for the magic parameter relation.
   * Ensures optimal bounce-back wall location for the two-relaxation-time
   * model. Default magic number is 3/16.
   */
  FloatType odd_mode_relaxation_rate(
      FloatType shear_relaxation,
      FloatType magic_number = FloatType{3} / FloatType{16}) const;
 
public:
  void set_collision_model(double kT, unsigned int seed) override;
  void set_collision_model(
      std::unique_ptr<LeesEdwardsPack> &&lees_edwards_pack) override;
  void check_lebc(unsigned int shear_direction,
                  unsigned int shear_plane_normal) const override;
 
public:
  // ---- Ghost Communication ----
 
  /**
   * @brief Perform all pending ghost layer updates.
   * Uses a lazy scheme: ghost communications are only executed when
   * they have been marked as pending by a preceding write operation.
   */
  void ghost_communication() override {
    if (m_pending_ghost_comm.any()) {
      assert(m_mpi_cart_comm_observer.is_valid());
      ghost_communication_boundary();
      ghost_communication_pdf();
      ghost_communication_laf();
      ghost_communication_vel();
    }
  }
 
  void ghost_communication_pdf() override {
    if (m_pending_ghost_comm.test(GhostComm::PDF)) {
      assert(m_mpi_cart_comm_observer.is_valid());
      m_pdf_communicator->communicate();
      if (has_lees_edwards_bc()) {
        auto const &blocks = get_lattice().get_blocks();
        apply_lees_edwards_pdf_interpolation(blocks);
      }
      m_pending_ghost_comm.reset(GhostComm::PDF);
    }
  }
 
  void ghost_communication_vel() override {
    if (m_pending_ghost_comm.test(GhostComm::VEL)) {
      assert(m_mpi_cart_comm_observer.is_valid());
      m_vel_communicator->communicate();
      if (has_lees_edwards_bc()) {
        auto const &blocks = get_lattice().get_blocks();
        apply_lees_edwards_vel_interpolation_and_shift(blocks);
      }
      m_pending_ghost_comm.reset(GhostComm::VEL);
    }
  }
 
  void ghost_communication_laf() override {
    if (m_pending_ghost_comm.test(GhostComm::LAF)) {
      assert(m_mpi_cart_comm_observer.is_valid());
      m_laf_communicator->communicate();
      if (has_lees_edwards_bc()) {
        auto const &blocks = get_lattice().get_blocks();
        apply_lees_edwards_last_applied_force_interpolation(blocks);
      }
      m_pending_ghost_comm.reset(GhostComm::LAF);
    }
  }
 
  void ghost_communication_boundary() {
    if (m_pending_ghost_comm.test(GhostComm::UBB)) {
      assert(m_mpi_cart_comm_observer.is_valid());
      m_boundary_communicator->communicate();
      m_pending_ghost_comm.reset(GhostComm::UBB);
    }
  }
 
  /** @brief Communicate all fields at once using the D3Q27 stencil. */
  void ghost_communication_full() {
    assert(m_mpi_cart_comm_observer.is_valid());
    m_full_communicator->communicate();
    if (has_lees_edwards_bc()) {
      apply_lees_edwards_interpolation();
    }
    m_pending_ghost_comm.reset(GhostComm::PDF);
    m_pending_ghost_comm.reset(GhostComm::VEL);
    m_pending_ghost_comm.reset(GhostComm::LAF);
  }
 
private:
  // ---- Lees-Edwards Boundary Conditions ----
 
  auto has_lees_edwards_bc() const {
    return std::holds_alternative<
        typename Kernels::StreamCollisionModelLeesEdwards>(*m_collision_model);
  }
 
  void apply_lees_edwards_pdf_interpolation(
      std::shared_ptr<BlockStorage> const &blocks) {
    for (auto &block : *blocks)
      (*m_lees_edwards_pdf_interpol_sweep)(&block);
  }
 
  void apply_lees_edwards_vel_interpolation_and_shift(
      std::shared_ptr<BlockStorage> const &blocks) {
    for (auto &block : *blocks)
      (*m_lees_edwards_vel_interpol_sweep)(&block);
  }
 
  void apply_lees_edwards_last_applied_force_interpolation(
      std::shared_ptr<BlockStorage> const &blocks) {
    for (auto &block : *blocks)
      (*m_lees_edwards_last_applied_force_interpol_sweep)(&block);
  }
 
public:
  void apply_lees_edwards_interpolation() {
    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);
  }
 
public:
  // ---- Node & Slice Accessors (by quantity) ----
 
  // Velocity
  std::optional<Utils::Vector3d>
  get_node_velocity(Utils::Vector3i const &node,
                    bool consider_ghosts = false) const override;
  bool set_node_velocity(Utils::Vector3i const &node,
                         Utils::Vector3d const &v) override;
  std::vector<double>
  get_slice_velocity(Utils::Vector3i const &lower_corner,
                     Utils::Vector3i const &upper_corner) const override;
  void set_slice_velocity(Utils::Vector3i const &lower_corner,
                          Utils::Vector3i const &upper_corner,
                          std::vector<double> const &velocity) override;
 
  // Density
  std::optional<double>
  get_node_density(Utils::Vector3i const &node,
                   bool consider_ghosts = false) const override;
  bool set_node_density(Utils::Vector3i const &node, double density) override;
  std::vector<double>
  get_slice_density(Utils::Vector3i const &lower_corner,
                    Utils::Vector3i const &upper_corner) const override;
  void set_slice_density(Utils::Vector3i const &lower_corner,
                         Utils::Vector3i const &upper_corner,
                         std::vector<double> const &density) override;
 
  // Population
  std::optional<std::vector<double>>
  get_node_population(Utils::Vector3i const &node,
                      bool consider_ghosts = false) const override;
  bool set_node_population(Utils::Vector3i const &node,
                           std::vector<double> const &population) override;
  std::vector<double>
  get_slice_population(Utils::Vector3i const &lower_corner,
                       Utils::Vector3i const &upper_corner) const override;
  void set_slice_population(Utils::Vector3i const &lower_corner,
                            Utils::Vector3i const &upper_corner,
                            std::vector<double> const &population) override;
 
  // Force
  std::optional<Utils::Vector3d>
  get_node_force_to_be_applied(Utils::Vector3i const &node) const override;
  std::optional<Utils::Vector3d>
  get_node_last_applied_force(Utils::Vector3i const &node,
                              bool consider_ghosts = false) const override;
  bool set_node_last_applied_force(Utils::Vector3i const &node,
                                   Utils::Vector3d const &force) override;
  std::vector<double> get_slice_last_applied_force(
      Utils::Vector3i const &lower_corner,
      Utils::Vector3i const &upper_corner) const override;
  void set_slice_last_applied_force(Utils::Vector3i const &lower_corner,
                                    Utils::Vector3i const &upper_corner,
                                    std::vector<double> const &force) override;
 
  // Pressure tensor
  std::optional<Utils::VectorXd<9>>
  get_node_pressure_tensor(Utils::Vector3i const &node) const override;
  std::vector<double>
  get_slice_pressure_tensor(Utils::Vector3i const &lower_corner,
                            Utils::Vector3i const &upper_corner) const override;
 
private:
  // ---- Interpolation (position-based access) ----
 
  /** @brief Return a B-spline interpolation kernel for force distribution. */
  auto make_force_interpolation_kernel() const;
  /** @brief Return a B-spline interpolation kernel for velocity readout. */
  auto make_velocity_interpolation_kernel() const;
  /** @brief Return a B-spline interpolation kernel for density readout. */
  auto make_density_interpolation_kernel() const;
 
public:
  std::function<bool(Utils::Vector3d const &)>
  make_lattice_position_checker(bool consider_points_in_halo) const override;
  bool add_force_at_pos(Utils::Vector3d const &pos,
                        Utils::Vector3d const &force) override;
  void add_forces_at_pos(std::vector<Utils::Vector3d> const &pos,
                         std::vector<Utils::Vector3d> const &forces) override;
  std::optional<Utils::Vector3d>
  get_velocity_at_pos(Utils::Vector3d const &pos,
                      bool consider_points_in_halo = false) const override;
  std::vector<Utils::Vector3d>
  get_velocities_at_pos(std::vector<Utils::Vector3d> const &pos) override;
  std::optional<double>
  get_density_at_pos(Utils::Vector3d const &pos,
                     bool consider_points_in_halo = false) const override;
  std::vector<double>
  get_densities_at_pos(std::vector<Utils::Vector3d> const &pos) override;
 
public:
  // ---- Boundary Handling ----
 
  void reset_boundary_handling(std::shared_ptr<BlockStorage> const &blocks) {
    auto const [lc, uc] = m_lattice->get_local_grid_range(true);
    m_boundary =
        std::make_shared<BoundaryModel>(blocks, m_pdf_field_id, m_flag_field_id,
                                        CellInterval{to_cell(lc), to_cell(uc)});
  }
 
  void on_boundary_add();
  void clear_boundaries() override;
  void reallocate_ubb_field() override;
  void
  update_boundary_from_shape(std::vector<int> const &raster_flat,
                             std::vector<double> const &data_flat) override;
  std::optional<Utils::Vector3d>
  get_node_velocity_at_boundary(Utils::Vector3i const &node,
                                bool consider_ghosts = false) const override;
  bool set_node_velocity_at_boundary(Utils::Vector3i const &node,
                                     Utils::Vector3d const &velocity) override;
  std::vector<std::optional<Utils::Vector3d>> get_slice_velocity_at_boundary(
      Utils::Vector3i const &lower_corner,
      Utils::Vector3i const &upper_corner) const override;
  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;
  std::optional<Utils::Vector3d>
  get_node_boundary_force(Utils::Vector3i const &node) const override;
  bool remove_node_from_boundary(Utils::Vector3i const &node) override;
  std::optional<bool>
  get_node_is_boundary(Utils::Vector3i const &node,
                       bool consider_ghosts = false) const override;
  std::vector<bool>
  get_slice_is_boundary(Utils::Vector3i const &lower_corner,
                        Utils::Vector3i const &upper_corner) const override;
  [[nodiscard]] Utils::Vector3d get_boundary_force_from_shape(
      std::vector<int> const &raster_flat) const override;
  [[nodiscard]] Utils::Vector3d get_boundary_force() const override;
 
private:
  [[nodiscard]] Utils::Vector3i flat_index_to_node(int index) const;
  [[nodiscard]] Utils::Vector3i get_neighbor_node(Utils::Vector3i const &node,
                                                  int dir) const;
 
public:
  // ---- Global Reductions & Physical Parameters ----
 
  // Global pressure tensor
  [[nodiscard]] Utils::VectorXd<9> get_pressure_tensor() const override {
    Matrix3<FloatType> tensor(FloatType{0});
    for (auto const &block : *get_lattice().get_blocks()) {
      auto pdf_field = block.template getData<PdfField>(m_pdf_field_id);
      tensor += lbm::accessor::PressureTensor::reduce(pdf_field, m_density);
    }
    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 {
    Vector3<FloatType> mom(FloatType{0});
    for (auto const &block : *get_lattice().get_blocks()) {
      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,
                                                    m_density);
    }
    return to_vector3d(mom);
  }
 
  // Global external force
  void set_external_force(Utils::Vector3d const &ext_force) override {
    m_reset_force->set_ext_force(zero_centered_to_lb(ext_force));
  }
 
  [[nodiscard]] Utils::Vector3d get_external_force() const noexcept override {
    return zero_centered_to_md(m_reset_force->get_ext_force());
  }
 
  void set_viscosity(double viscosity) override {
    m_viscosity = FloatType_c(viscosity);
  }
 
  [[nodiscard]] double get_viscosity() const noexcept override {
    return static_cast<double>(m_viscosity);
  }
 
  [[nodiscard]] double get_density() const noexcept override {
    return static_cast<double>(m_density);
  }
 
  [[nodiscard]] double get_kT() const noexcept override {
    return static_cast<double>(m_kT);
  }
 
  [[nodiscard]] unsigned int get_seed() const noexcept override {
    return m_seed;
  }
 
  [[nodiscard]] std::optional<uint64_t> get_rng_state() const override {
    auto const cm =
        std::get_if<typename Kernels::StreamCollisionModelThermalized>(
            &*m_collision_model);
    if (!cm or m_kT == 0.) {
      return std::nullopt;
    }
    return {static_cast<uint64_t>(cm->getTime_step())};
  }
 
  void set_rng_state(uint64_t counter) override {
    auto const cm =
        std::get_if<typename Kernels::StreamCollisionModelThermalized>(
            &*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->setTime_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;
  }
 
  /**
   * @brief Correction factor for off-diagonal pressure tensor elements.
   * Compensates for the viscosity-dependent error in the non-equilibrium
   * stress: factor = nu / (nu + 1/6).
   */
  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;
    }
  }
 
  void pressure_tensor_correction(std::span<FloatType, 9ul> tensor) const {
    auto const revert_factor = pressure_tensor_correction_factor();
    for (auto const i : {1u, 2u, 3u, 5u, 6u, 7u}) {
      tensor[i] *= revert_factor;
    }
  }
 
protected:
  /**
   * @brief Scale data by a conversion factor (in-place).
   * Used for zero-centered density representation: LB internally stores
   * density fluctuations around zero, while the user interface uses
   * absolute densities. The conversion factors @ref m_zc_to_md and
   * @ref m_zc_to_lb translate between these representations.
   */
  template <typename T>
  void zero_centered_transform_impl(T &data, auto const factor) const {
    if constexpr (std::is_arithmetic_v<T>) {
      static_assert(std::is_floating_point_v<T>);
      data *= static_cast<T>(factor);
    } else {
      auto const coef = static_cast<typename T::value_type>(factor);
      std::transform(std::begin(data), std::end(data), std::begin(data),
                     [coef](auto value) { return value * coef; });
    }
  }
 
  void zero_centered_to_lb_in_place(auto &data) const {
    zero_centered_transform_impl(data, m_zc_to_lb);
  }
 
  void zero_centered_to_md_in_place(auto &data) const {
    zero_centered_transform_impl(data, m_zc_to_md);
  }
 
  auto zero_centered_to_lb(auto const &data) const {
    auto transformed_data = data;
    zero_centered_to_lb_in_place(transformed_data);
    return transformed_data;
  }
 
  auto zero_centered_to_md(auto const &data) const {
    auto transformed_data = data;
    zero_centered_to_md_in_place(transformed_data);
    return transformed_data;
  }
 
public:
  // ---- File I/O ----
 
  void register_vtk_field_filters(walberla::vtk::VTKOutput &vtk_obj) override {
    field::FlagFieldCellFilter<FlagField> fluid_filter(m_flag_field_id);
    fluid_filter.addFlag(Boundary_flag);
    vtk_obj.addCellExclusionFilter(fluid_filter);
  }
 
  void register_vtk_field_writers(walberla::vtk::VTKOutput &vtk_obj,
                                  LatticeModel::units_map const &units,
                                  int flag_observables) override;
 
protected:
  // ---- Private Infrastructure Helpers ----
 
  /**
   * @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
      constexpr auto alignment = field::SIMDAlignment();
      using value_type = 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__) and defined(WALBERLA_BUILD_WITH_CUDA)
    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) {
          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) {
          auto field = block.template getData<GPUField>(field_id);
          lbm::accessor::Population::initialize(
              field, std::array<FloatType, Stencil::Size>{});
        }
      }
      return field_id;
    }
#endif
  }
 
  /**
   * @brief Set up D3Q27 communicators for full ghost layer updates.
   * Creates per-field communicators (PDF, velocity, last-applied force)
   * as well as a combined communicator and the boundary communicator.
   */
  void setup_full_communicator() {
    auto const &blocks = m_lattice->get_blocks();
 
    m_full_communicator = std::make_shared<RegularFullCommunicator>(blocks);
    m_full_communicator->addPackInfo(
        std::make_shared<PackInfo<PdfField>>(m_pdf_field_id));
    m_full_communicator->addPackInfo(
        std::make_shared<PackInfo<VectorField>>(m_last_applied_force_field_id));
    m_full_communicator->addPackInfo(
        std::make_shared<PackInfo<VectorField>>(m_velocity_field_id));
 
    m_pdf_communicator = std::make_shared<RegularFullCommunicator>(blocks);
    m_vel_communicator = std::make_shared<RegularFullCommunicator>(blocks);
    m_laf_communicator = std::make_shared<RegularFullCommunicator>(blocks);
    m_pdf_communicator->addPackInfo(
        std::make_shared<PackInfo<PdfField>>(m_pdf_field_id));
    m_vel_communicator->addPackInfo(
        std::make_shared<PackInfo<VectorField>>(m_velocity_field_id));
    m_laf_communicator->addPackInfo(
        std::make_shared<PackInfo<VectorField>>(m_last_applied_force_field_id));
 
    m_boundary_communicator =
        std::make_shared<BoundaryFullCommunicator>(blocks);
    m_boundary_communicator->addPackInfo(
        std::make_shared<field::communication::BoundaryFlagPackInfo<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);
  }
 
  /**
   * @brief Set up the communicator used during integration.
   * Uses optimized streaming pack info when neither boundaries nor
   * Lees-Edwards boundary conditions are active; falls back to the
   * generic pack info otherwise.
   */
  void setup_streaming_communicator() {
    auto const setup = [this]<typename PackInfoPdf, typename PackInfoVec>() {
      auto const &blocks = m_lattice->get_blocks();
      m_pdf_streaming_communicator =
          std::make_shared<PDFStreamingCommunicator>(blocks);
      m_pdf_streaming_communicator->addPackInfo(
          std::make_shared<PackInfoPdf>(m_pdf_field_id));
      m_pdf_streaming_communicator->addPackInfo(
          std::make_shared<PackInfoVec>(m_last_applied_force_field_id));
    };
    using FieldTrait = FieldTrait<FloatType, Stencil, Architecture>;
    using PackInfoPdf = FieldTrait::PackInfoStreamingPdf;
    using PackInfoVec = FieldTrait::PackInfoStreamingVec;
    if (m_has_boundaries or (m_collision_model and has_lees_edwards_bc())) {
      setup.template operator()<PackInfo<PdfField>, PackInfoVec>();
    } else {
      setup.template operator()<PackInfoPdf, PackInfoVec>();
    }
  }
};
 
} // namespace walberla
 
// Out-of-class template method definitions
#include "LBBoundaryAccess.impl.hpp"
#include "LBCollisionSetup.impl.hpp"
#include "LBInterpolation.impl.hpp"
#include "LBNodeAccess.impl.hpp"
#include "LBSliceAccess.impl.hpp"
#include "LBVTK.impl.hpp"