Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve LB GPU and thermostats API #5014

Merged
merged 3 commits into from
Nov 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion doc/tutorials/electrokinetics/electrokinetics.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -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)"
]
Expand Down
56 changes: 41 additions & 15 deletions maintainer/walberla_kernels/templates/FieldAccessors.tmpl.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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 )
Expand All @@ -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 %}
}
}
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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


Expand Down
2 changes: 2 additions & 0 deletions maintainer/walberla_kernels/templates/FieldAccessors.tmpl.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
21 changes: 21 additions & 0 deletions maintainer/walberla_kernels/templates/FieldAccessors.tmpl.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 0 additions & 1 deletion samples/lb_profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down
27 changes: 16 additions & 11 deletions src/core/lb/particle_coupling.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
3 changes: 3 additions & 0 deletions src/core/thermostat.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
70 changes: 58 additions & 12 deletions src/script_interface/thermostat/thermostat.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
#include <cassert>
#include <limits>
#include <memory>
#include <span>
#include <stdexcept>
#include <string>

Expand Down Expand Up @@ -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();
}

Expand All @@ -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);
Expand All @@ -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) {
Expand Down Expand Up @@ -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);
}
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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({
Expand All @@ -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
Expand All @@ -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({
Expand All @@ -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
Expand All @@ -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({
Expand Down Expand Up @@ -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

Expand All @@ -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; }
};
Expand All @@ -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; }
};
Expand All @@ -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; }
};
Expand Down
Loading