Skip to content

Commit

Permalink
core: Use zip iterators in the P3M algorithm
Browse files Browse the repository at this point in the history
  • Loading branch information
@jhossbach @jngrad
jhossbach authored and jngrad committed Dec 22, 2023
1 parent cc3a319 commit 7d6434b
Show file tree
Hide file tree
Showing 4 changed files with 62 additions and 42 deletions.
14 changes: 10 additions & 4 deletions src/core/ParticlePropertyIterator.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,19 +42,25 @@ auto create_transform_range(ParticleRange const &particles, Kernel kernel) {
}
} // namespace detail

auto unfolded_pos_range(ParticleRange const &particles,
BoxGeometry const &box) {
inline auto unfolded_pos_range(ParticleRange const &particles,
BoxGeometry const &box) {
auto return_unfolded_pos = [&box](Particle &p) {
return ::detail::unfolded_position(p.pos(), p.image_box(), box.length());
};
return detail::create_transform_range(particles, return_unfolded_pos);
}

auto charge_range(ParticleRange const &particles) {
inline auto pos_range(ParticleRange const &particles) {
auto return_pos = [](Particle &p) -> Utils::Vector3d & { return p.pos(); };
return detail::create_transform_range(particles, return_pos);
}

inline auto charge_range(ParticleRange const &particles) {
auto return_charge = [](Particle &p) -> double & { return p.q(); };
return detail::create_transform_range(particles, return_charge);
}
auto force_range(ParticleRange const &particles) {

inline auto force_range(ParticleRange const &particles) {
auto return_force = [](Particle &p) -> Utils::Vector3d & {
return p.force();
};
Expand Down
68 changes: 40 additions & 28 deletions src/core/electrostatics/elc.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@

#include "BoxGeometry.hpp"
#include "Particle.hpp"
#include "ParticlePropertyIterator.hpp"
#include "ParticleRange.hpp"
#include "cell_system/CellStructure.hpp"
#include "communication.hpp"
Expand All @@ -43,6 +44,7 @@
#include <utils/math/sqr.hpp>

#include <boost/mpi/collectives/all_reduce.hpp>
#include <boost/range/combine.hpp>

#include <algorithm>
#include <cassert>
Expand Down Expand Up @@ -1118,59 +1120,62 @@ static void p3m_assign_image_charge(elc_data const &elc, CoulombP3M &p3m,
}
}

template <ChargeProtocol protocol>
template <ChargeProtocol protocol, typename combined_ranges>
void charge_assign(elc_data const &elc, CoulombP3M &solver,
ParticleRange const &particles) {
combined_ranges const &p_q_pos_range) {
if (protocol == ChargeProtocol::BOTH or protocol == ChargeProtocol::IMAGE) {
solver.p3m.inter_weights.reset(solver.p3m.params.cao);
}
/* prepare local FFT mesh */
for (int i = 0; i < solver.p3m.local_mesh.size; i++)
solver.p3m.rs_mesh[i] = 0.;

for (auto const &p : particles) {
if (p.q() != 0.) {
for (auto zipped : p_q_pos_range) {
auto const p_q = boost::get<0>(zipped);
auto const &p_pos = boost::get<1>(zipped);
if (p_q != 0.) {
if (protocol == ChargeProtocol::BOTH or
protocol == ChargeProtocol::REAL) {
solver.assign_charge(p.q(), p.pos(), solver.p3m.inter_weights);
solver.assign_charge(p_q, p_pos, solver.p3m.inter_weights);
}
if (protocol == ChargeProtocol::BOTH or
protocol == ChargeProtocol::IMAGE) {
p3m_assign_image_charge(elc, solver, p.q(), p.pos());
p3m_assign_image_charge(elc, solver, p_q, p_pos);
}
}
}
}

template <ChargeProtocol protocol>
template <ChargeProtocol protocol, typename combined_range>
void modify_p3m_sums(elc_data const &elc, CoulombP3M &solver,
ParticleRange const &particles) {
combined_range const &p_q_pos_range) {

Utils::Vector3d node_sums{};
for (auto const &p : particles) {
auto const q = p.q();
if (q != 0.) {
auto const z = p.pos()[2];
for (auto zipped : p_q_pos_range) {
auto const p_q = boost::get<0>(zipped);
auto const &p_pos = boost::get<1>(zipped);
if (p_q != 0.) {
auto const p_z = p_pos[2];

if (protocol == ChargeProtocol::BOTH or
protocol == ChargeProtocol::REAL) {
node_sums[0] += 1.;
node_sums[1] += Utils::sqr(q);
node_sums[2] += q;
node_sums[1] += Utils::sqr(p_q);
node_sums[2] += p_q;
}

if (protocol == ChargeProtocol::BOTH or
protocol == ChargeProtocol::IMAGE) {
if (z < elc.space_layer) {
if (p_z < elc.space_layer) {
node_sums[0] += 1.;
node_sums[1] += Utils::sqr(elc.delta_mid_bot * q);
node_sums[2] += elc.delta_mid_bot * q;
node_sums[1] += Utils::sqr(elc.delta_mid_bot * p_q);
node_sums[2] += elc.delta_mid_bot * p_q;
}

if (z > (elc.box_h - elc.space_layer)) {
if (p_z > (elc.box_h - elc.space_layer)) {
node_sums[0] += 1.;
node_sums[1] += Utils::sqr(elc.delta_mid_top * q);
node_sums[2] += elc.delta_mid_top * q;
node_sums[1] += Utils::sqr(elc.delta_mid_top * p_q);
node_sums[2] += elc.delta_mid_top * p_q;
}
}
}
Expand All @@ -1190,6 +1195,10 @@ double ElectrostaticLayerCorrection::long_range_energy(
auto &solver = *solver_ptr;
auto const &box_geo = *get_system().box_geo;

auto p_q_range = ParticlePropertyRange::charge_range(particles);
auto p_pos_range = ParticlePropertyRange::pos_range(particles);
auto p_q_pos_range = boost::combine(p_q_range, p_pos_range);

// assign the original charges (they may not have been assigned yet)
solver.charge_assign(particles);

Expand All @@ -1203,17 +1212,17 @@ double ElectrostaticLayerCorrection::long_range_energy(
0.5 * elc.dielectric_layers_self_energy(solver, box_geo, particles);

// assign both original and image charges
charge_assign<ChargeProtocol::BOTH>(elc, solver, particles);
modify_p3m_sums<ChargeProtocol::BOTH>(elc, solver, particles);
charge_assign<ChargeProtocol::BOTH>(elc, solver, p_q_pos_range);
modify_p3m_sums<ChargeProtocol::BOTH>(elc, solver, p_q_pos_range);
energy += 0.5 * solver.long_range_energy(particles);

// assign only the image charges now
charge_assign<ChargeProtocol::IMAGE>(elc, solver, particles);
modify_p3m_sums<ChargeProtocol::IMAGE>(elc, solver, particles);
charge_assign<ChargeProtocol::IMAGE>(elc, solver, p_q_pos_range);
modify_p3m_sums<ChargeProtocol::IMAGE>(elc, solver, p_q_pos_range);
energy -= 0.5 * solver.long_range_energy(particles);

// restore modified sums
modify_p3m_sums<ChargeProtocol::REAL>(elc, solver, particles);
modify_p3m_sums<ChargeProtocol::REAL>(elc, solver, p_q_pos_range);

return energy;
},
Expand All @@ -1226,17 +1235,20 @@ void ElectrostaticLayerCorrection::add_long_range_forces(
std::visit(
[this, &particles](auto const &solver_ptr) {
auto &solver = *solver_ptr;
auto p_q_range = ParticlePropertyRange::charge_range(particles);
auto p_pos_range = ParticlePropertyRange::pos_range(particles);
auto p_q_pos_range = boost::combine(p_q_range, p_pos_range);
if (elc.dielectric_contrast_on) {
auto const &box_geo = *get_system().box_geo;
modify_p3m_sums<ChargeProtocol::BOTH>(elc, solver, particles);
charge_assign<ChargeProtocol::BOTH>(elc, solver, particles);
modify_p3m_sums<ChargeProtocol::BOTH>(elc, solver, p_q_pos_range);
charge_assign<ChargeProtocol::BOTH>(elc, solver, p_q_pos_range);
elc.dielectric_layers_self_forces(solver, box_geo, particles);
} else {
solver.charge_assign(particles);
}
solver.add_long_range_forces(particles);
if (elc.dielectric_contrast_on) {
modify_p3m_sums<ChargeProtocol::REAL>(elc, solver, particles);
modify_p3m_sums<ChargeProtocol::REAL>(elc, solver, p_q_pos_range);
}
},
base_solver);
Expand Down
18 changes: 12 additions & 6 deletions src/core/electrostatics/p3m.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -323,10 +323,13 @@ template <int cao> struct AssignCharge {
[q, &p3m](int ind, double w) { p3m.rs_mesh[ind] += w * q; });
}

void operator()(p3m_data_struct &p3m, ParticleRange const &particles) {
for (auto &p : particles) {
if (p.q() != 0.0) {
this->operator()(p3m, p.q(), p.pos(), p3m.inter_weights);
template <typename combined_ranges>
void operator()(p3m_data_struct &p3m, combined_ranges const &p_q_pos_range) {
for (auto zipped : p_q_pos_range) {
auto const p_q = boost::get<0>(zipped);
auto const &p_pos = boost::get<1>(zipped);
if (p_q != 0.0) {
this->operator()(p3m, p_q, p_pos, p3m.inter_weights);
}
}
}
Expand All @@ -340,8 +343,11 @@ void CoulombP3M::charge_assign(ParticleRange const &particles) {
for (int i = 0; i < p3m.local_mesh.size; i++)
p3m.rs_mesh[i] = 0.0;

Utils::integral_parameter<int, AssignCharge, 1, 7>(p3m.params.cao, p3m,
particles);
auto p_q_range = ParticlePropertyRange::charge_range(particles);
auto p_pos_range = ParticlePropertyRange::pos_range(particles);

Utils::integral_parameter<int, AssignCharge, 1, 7>(
p3m.params.cao, p3m, boost::combine(p_q_range, p_pos_range));
}

void CoulombP3M::assign_charge(double q, Utils::Vector3d const &real_pos,
Expand Down
4 changes: 0 additions & 4 deletions src/core/electrostatics/p3m.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,10 +99,6 @@ struct CoulombP3M : public Coulomb::Actor<CoulombP3M> {

bool is_tuned() const { return m_is_tuned; }

/** Compute the k-space part of forces and energies. */
double kernel(bool force_flag, bool energy_flag,
ParticleRange const &particles);

/** @brief Recalculate all derived parameters. */
void init();
void on_activation() {
Expand Down

0 comments on commit 7d6434b

Please sign in to comment.