Improve LB GPU and thermostats API (#5014)

Description of changes:
- LB GPU now supports VTK output
- LB GPU total pressure tensor and total momentum are now significantly faster
- thermostat parameter `gamma` is now required
   - fixes a regression introduced in the last thermostat API refactoring

doc/tutorials/electrokinetics/electrokinetics.ipynb

@@ -1042,7 +1042,7 @@
     "                lattice=lattice, density=DENSITY_FLUID, kinematic_viscosity=VISCOSITY_KINEMATIC,\n",
     "                tau=TAU, ext_force_density=EXT_FORCE_DENSITY, kT=KT, seed=42)\n",
     "system.lb = lbf\n",
-    "system.thermostat.set_lb(LB_fluid=lbf, seed=42)\n",
+    "system.thermostat.set_lb(LB_fluid=lbf, seed=42, gamma=0.)\n",
     "eksolver = espressomd.electrokinetics.EKNone(lattice=lattice)\n",
     "system.ekcontainer = espressomd.electrokinetics.EKContainer(tau=TAU, solver=eksolver)"
    ]

maintainer/walberla_kernels/templates/FieldAccessors.tmpl.cu

@@ -899,7 +899,7 @@ namespace Force {
 namespace MomentumDensity
 {
 // LCOV_EXCL_START
-    __global__ void kernel_sum(
+    __global__ void kernel_get(
         gpu::FieldAccessor< {{dtype}} > pdf,
         gpu::FieldAccessor< {{dtype}} > force,
         {{dtype}} * RESTRICT out )
@@ -914,7 +914,7 @@ namespace MomentumDensity
             {% endfor -%}
             {{momentum_density_getter | substitute_force_getter_cu | indent(8) }}
             {% for i in range(D) -%}
-                out[{{i}}u] += md_{{i}};
+                out[{{i}}u] = md_{{i}};
             {% endfor %}
         }
     }
@@ -924,19 +924,22 @@ namespace MomentumDensity
         gpu::GPUField< {{dtype}} > const * pdf_field,
         gpu::GPUField< {{dtype}} > const * force_field )
     {
-        thrust::device_vector< {{dtype}} > dev_data({{D}}u, {{dtype}} {0});
+        auto const ci = pdf_field->xyzSize();
+        thrust::device_vector< {{dtype}} > dev_data({{D}}u * ci.numCells());
         auto const dev_data_ptr = thrust::raw_pointer_cast(dev_data.data());
-        WALBERLA_FOR_ALL_CELLS_XYZ(pdf_field, {
-            Cell cell(x, y, z);
-            CellInterval ci ( cell, cell );
-            auto kernel = gpu::make_kernel( kernel_sum );
-            kernel.addFieldIndexingParam( gpu::FieldIndexing< {{dtype}} >::interval( *pdf_field, ci ) );
-            kernel.addFieldIndexingParam( gpu::FieldIndexing< {{dtype}} >::interval( *force_field, ci ) );
-            kernel.addParam( dev_data_ptr );
-            kernel();
-        });
+        auto kernel = gpu::make_kernel( kernel_get );
+        kernel.addFieldIndexingParam( gpu::FieldIndexing< {{dtype}} >::interval( *pdf_field, ci) );
+        kernel.addFieldIndexingParam( gpu::FieldIndexing< {{dtype}} >::interval( *force_field, ci ) );
+        kernel.addParam( dev_data_ptr );
+        kernel();
+        std::vector< {{dtype}} > out(dev_data.size());
+        thrust::copy(dev_data.begin(), dev_data.end(), out.data());
         Vector{{D}}< {{dtype}} > mom({{dtype}} {0});
-        thrust::copy(dev_data.begin(), dev_data.begin() + {{D}}u, mom.data());
+        for (auto i = uint_t{ 0u }; i < {{D}}u * ci.numCells(); i += {{D}}u) {
+            {% for j in range(D) -%}
+                mom[{{j}}u] += out[i + {{j}}u];
+            {% endfor -%}
+        }
         return mom;
     }
 } // namespace MomentumDensity
@@ -977,7 +980,7 @@ namespace PressureTensor
         Matrix{{D}}< {{dtype}} > out;
         thrust::copy(dev_data.begin(), dev_data.end(), out.data());
         return out;
-   }
+    }
 
     std::vector< {{dtype}} > get(
         gpu::GPUField< {{dtype}} > const * pdf_field,
@@ -992,7 +995,30 @@ namespace PressureTensor
         std::vector< {{dtype}} > out(dev_data.size());
         thrust::copy(dev_data.begin(), dev_data.end(), out.data());
         return out;
-   }
+    }
+
+    Matrix{{D}}< {{dtype}} > reduce(
+        gpu::GPUField< {{dtype}} > const * pdf_field)
+    {
+        auto const ci = pdf_field->xyzSize();
+        thrust::device_vector< {{dtype}} > dev_data({{D**2}}u * ci.numCells());
+        auto const dev_data_ptr = thrust::raw_pointer_cast(dev_data.data());
+        auto kernel = gpu::make_kernel( kernel_get );
+        kernel.addFieldIndexingParam( gpu::FieldIndexing< {{dtype}} >::interval( *pdf_field, ci ) );
+        kernel.addParam( dev_data_ptr );
+        kernel();
+        std::vector< {{dtype}} > out(dev_data.size());
+        thrust::copy(dev_data.begin(), dev_data.end(), out.data());
+        Matrix{{D}}< {{dtype}} > pressureTensor({{dtype}} {0});
+        for (auto i = uint_t{ 0u }; i < {{D**2}}u * ci.numCells(); i += {{D**2}}u) {
+            {% for i in range(D) -%}
+                {% for j in range(D) -%}
+                    pressureTensor[{{i*D+j}}u] += out[i + {{i*D+j}}u];
+                {% endfor %}
+            {% endfor %}
+        }
+        return pressureTensor;
+    }
 } // namespace PressureTensor
 
 

maintainer/walberla_kernels/templates/FieldAccessors.tmpl.cuh

@@ -213,6 +213,8 @@ namespace PressureTensor {
     std::vector< {{dtype}} >
     get( gpu::GPUField< {{dtype}} > const * pdf_field,
          CellInterval const & ci );
+    Matrix{{D}}< {{dtype}} >
+    reduce( gpu::GPUField< {{dtype}} > const * pdf_field );
 } // namespace PressureTensor
 
 } // namespace accessor

maintainer/walberla_kernels/templates/FieldAccessors.tmpl.h

@@ -641,6 +641,27 @@ namespace PressureTensor
         }
         return out;
     }
+
+    inline auto
+    reduce( GhostLayerField< {{dtype}}, uint_t{ {{Q}}u } > const * pdf_field)
+    {
+        Matrix{{D}}< {{dtype}} > pressureTensor({{dtype}} {0});
+        WALBERLA_FOR_ALL_CELLS_XYZ(pdf_field, {
+            const {{dtype}} & xyz0 = pdf_field->get(x, y, z, uint_t{ 0u });
+            {% for i in range(Q) -%}
+                const {{dtype}} f_{{i}} = pdf_field->getF( &xyz0, uint_t{ {{i}}u });
+            {% endfor -%}
+
+            {{second_momentum_getter | indent(8) }}
+
+            {% for i in range(D) -%}
+                {% for j in range(D) -%}
+                    pressureTensor[{{i*D+j}}u] += p_{{i*D+j}};
+                {% endfor %}
+            {% endfor %}
+        });
+        return pressureTensor;
+    }
 } // namespace PressureTensor
 
 } // namespace accessor

samples/lb_profile.py

@@ -45,7 +45,6 @@
     agrid=1.0, density=1.0, kinematic_viscosity=1.0, tau=0.01,
     ext_force_density=[0, 0, 0.15], kT=0.0)
 system.lb = lb_fluid
-system.thermostat.set_lb(LB_fluid=lb_fluid, seed=23)
 ctp = espressomd.math.CylindricalTransformationParameters(
     center=[5.0, 5.0, 0.0],
     axis=[0, 0, 1],

src/core/lb/particle_coupling.cpp

@@ -291,18 +291,23 @@ void ParticleCoupling::kernel(std::vector<Particle *> const &particles) {
 #endif
     Utils::Vector3d force_on_particle = {};
     if (coupling_mode == particle_force) {
-      auto &v_fluid = *it_interpolated_velocities;
-      if (m_box_geo.type() == BoxType::LEES_EDWARDS) {
-        // Account for the case where the interpolated velocity has been read
-        // from a ghost of the particle across the LE boundary (or vice verssa)
-        // Then the particle velocity is shifted by +,- the LE shear velocity
-        auto const vel_correction = lees_edwards_vel_shift(
-            *it_positions_velocity_coupling, p.pos(), m_box_geo);
-        v_fluid += vel_correction;
+#ifndef THERMOSTAT_PER_PARTICLE
+      if (m_thermostat.gamma > 0.)
+#endif
+      {
+        auto &v_fluid = *it_interpolated_velocities;
+        if (m_box_geo.type() == BoxType::LEES_EDWARDS) {
+          // Account for the case where the interpolated velocity has been read
+          // from a ghost of the particle across the LE boundary (or vice versa)
+          // Then the particle velocity is shifted by +,- the LE shear velocity
+          auto const vel_correction = lees_edwards_vel_shift(
+              *it_positions_velocity_coupling, p.pos(), m_box_geo);
+          v_fluid += vel_correction;
+        }
+        auto const drag_force = lb_drag_force(p, m_thermostat.gamma, v_fluid);
+        auto const random_force = get_noise_term(p);
+        force_on_particle = drag_force + random_force;
       }
-      auto const drag_force = lb_drag_force(p, m_thermostat.gamma, v_fluid);
-      auto const random_force = get_noise_term(p);
-      force_on_particle = drag_force + random_force;
       ++it_interpolated_velocities;
       ++it_positions_velocity_coupling;
     }

src/core/thermostat.hpp

@@ -97,6 +97,9 @@ inline bool are_kT_equal(double old_kT, double new_kT) {
   }
   auto const large_kT = (old_kT > new_kT) ? old_kT : new_kT;
   auto const small_kT = (old_kT > new_kT) ? new_kT : old_kT;
+  if (small_kT == 0.) {
+    return false;
+  }
   return (large_kT / small_kT - 1. < relative_tolerance);
 }
 } // namespace Thermostat

src/script_interface/thermostat/thermostat.hpp

@@ -33,6 +33,7 @@
 #include <cassert>
 #include <limits>
 #include <memory>
+#include <span>
 #include <stdexcept>
 #include <string>
 
@@ -92,7 +93,7 @@ class Interface : public AutoParameters<Interface<CoreClass>, System::Leaf> {
   }
 
   virtual bool invalid_rng_state(VariantMap const &params) const {
-    return (not params.count("seed") or is_none(params.at("seed"))) and
+    return (not params.contains("seed") or is_none(params.at("seed"))) and
            is_seed_required();
   }
 
@@ -101,14 +102,15 @@ class Interface : public AutoParameters<Interface<CoreClass>, System::Leaf> {
   get_member_handle(::Thermostat::Thermostat &thermostat) = 0;
 
   void set_new_parameters(VariantMap const &params) {
-    if (params.count("__check_rng_state") and invalid_rng_state(params)) {
-      context()->parallel_try_catch([]() {
+    context()->parallel_try_catch([&]() {
+      if (params.contains("__check_rng_state") and invalid_rng_state(params)) {
         throw std::invalid_argument("Parameter 'seed' is needed on first "
                                     "activation of the thermostat");
-      });
-    }
+      }
+      check_required_parameters(params);
+    });
     for (auto const &key : get_parameter_insertion_order()) {
-      if (params.count(key)) {
+      if (params.contains(key)) {
         auto const &v = params.at(key);
         if (key == "is_active") {
           is_active = get_value<bool>(v);
@@ -119,6 +121,17 @@ class Interface : public AutoParameters<Interface<CoreClass>, System::Leaf> {
     }
   }
 
+  virtual std::span<std::string_view const> get_required_parameters() const = 0;
+
+  void check_required_parameters(VariantMap const &params) const {
+    for (auto const &required : get_required_parameters()) {
+      auto name = std::string(required);
+      if (not params.contains(name)) {
+        throw std::runtime_error("Parameter '" + name + "' is missing");
+      }
+    }
+  }
+
 protected:
   template <typename T>
   auto make_autoparameter(T CoreThermostat::*member, char const *name) {
@@ -169,7 +182,7 @@ class Interface : public AutoParameters<Interface<CoreClass>, System::Leaf> {
     auto params = parameters;
     if (not is_seed_required()) {
       for (auto key : {std::string("seed"), std::string("philox_counter")}) {
-        if (params.count(key) == 0ul) {
+        if (not params.contains(key)) {
           params[key] = get_parameter(key);
         }
       }
@@ -216,7 +229,7 @@ class Interface : public AutoParameters<Interface<CoreClass>, System::Leaf> {
   }
 
   virtual std::optional<double> extract_kT(VariantMap const &params) const {
-    if (params.count("kT")) {
+    if (params.contains("kT")) {
       auto const value = get_value<double>(params, "kT");
       sanity_checks_positive(value, "kT");
       return value;
@@ -301,6 +314,11 @@ class Langevin : public Interface<::LangevinThermostat> {
     return thermostat.langevin;
   }
 
+  std::span<std::string_view const> get_required_parameters() const override {
+    static constexpr std::array names{std::string_view("gamma")};
+    return names;
+  }
+
 public:
   Langevin() {
     add_parameters({
@@ -319,7 +337,7 @@ class Langevin : public Interface<::LangevinThermostat> {
         Interface<::LangevinThermostat>::extend_parameters(parameters);
 #ifdef ROTATION
     // If gamma_rotation is not set explicitly, use the translational one.
-    if (params.count("gamma_rotation") == 0ul and params.count("gamma")) {
+    if (not params.contains("gamma_rotation") and params.contains("gamma")) {
       params["gamma_rotation"] = params.at("gamma");
     }
 #endif // ROTATION
@@ -333,6 +351,11 @@ class Brownian : public Interface<::BrownianThermostat> {
     return thermostat.brownian;
   }
 
+  std::span<std::string_view const> get_required_parameters() const override {
+    static constexpr std::array names{std::string_view("gamma")};
+    return names;
+  }
+
 public:
   Brownian() {
     add_parameters({
@@ -351,7 +374,7 @@ class Brownian : public Interface<::BrownianThermostat> {
         Interface<::BrownianThermostat>::extend_parameters(parameters);
 #ifdef ROTATION
     // If gamma_rotation is not set explicitly, use the translational one.
-    if (params.count("gamma_rotation") == 0ul and params.count("gamma")) {
+    if (not params.contains("gamma_rotation") and params.contains("gamma")) {
       params["gamma_rotation"] = params.at("gamma");
     }
 #endif // ROTATION
@@ -366,6 +389,12 @@ class IsotropicNpt : public Interface<::IsotropicNptThermostat> {
     return thermostat.npt_iso;
   }
 
+  std::span<std::string_view const> get_required_parameters() const override {
+    static constexpr std::array names{std::string_view("gamma0"),
+                                      std::string_view("gammav")};
+    return names;
+  }
+
 public:
   IsotropicNpt() {
     add_parameters({
@@ -418,10 +447,15 @@ class LBThermostat : public Interface<::LBThermostat> {
 
 protected:
   bool invalid_rng_state(VariantMap const &params) const override {
-    return (not params.count("seed") or is_none(params.at("seed"))) and
-           params.count("__global_kT") and is_seed_required() and
+    return (not params.contains("seed") or is_none(params.at("seed"))) and
+           params.contains("__global_kT") and is_seed_required() and
            get_value<double>(params, "__global_kT") != 0.;
   }
+
+  std::span<std::string_view const> get_required_parameters() const override {
+    static constexpr std::array names{std::string_view("gamma")};
+    return names;
+  }
 };
 #endif // WALBERLA
 
@@ -432,6 +466,10 @@ class DPDThermostat : public Interface<::DPDThermostat> {
     return thermostat.dpd;
   }
 
+  std::span<std::string_view const> get_required_parameters() const override {
+    return {};
+  }
+
 public:
   ::ThermostatFlags get_thermo_flag() const final { return THERMO_DPD; }
 };
@@ -444,6 +482,10 @@ class Stokesian : public Interface<::StokesianThermostat> {
     return thermostat.stokesian;
   }
 
+  std::span<std::string_view const> get_required_parameters() const override {
+    return {};
+  }
+
 public:
   ::ThermostatFlags get_thermo_flag() const final { return THERMO_SD; }
 };
@@ -455,6 +497,10 @@ class ThermalizedBond : public Interface<::ThermalizedBondThermostat> {
     return thermostat.thermalized_bond;
   }
 
+  std::span<std::string_view const> get_required_parameters() const override {
+    return {};
+  }
+
 public:
   ::ThermostatFlags get_thermo_flag() const final { return THERMO_BOND; }
 };