Skip to content

Commit

Permalink
remove particlerange from charge_assign
Browse files Browse the repository at this point in the history
  • Loading branch information
@jhossbach
jhossbach committed Dec 8, 2023
1 parent 3a2b6dd commit 36b27fd
Show file tree
Hide file tree
Showing 3 changed files with 158 additions and 83 deletions.
78 changes: 48 additions & 30 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,8 +1195,15 @@ double ElectrostaticLayerCorrection::long_range_energy(
auto &solver = *solver_ptr;
auto const &box_geo = *get_system().box_geo;

// TODO: Move construction up
//
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);
solver.charge_assign(p_q_pos_range);

if (!elc.dielectric_contrast_on) {
return solver.long_range_energy(particles);
Expand All @@ -1203,17 +1215,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 @@ -1225,18 +1237,24 @@ void ElectrostaticLayerCorrection::add_long_range_forces(
ParticleRange const &particles) const {
std::visit(
[this, &particles](auto const &solver_ptr) {
// TODO: Move construction up
//
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);
auto &solver = *solver_ptr;
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.charge_assign(p_q_pos_range);
}
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
109 changes: 57 additions & 52 deletions src/core/electrostatics/p3m.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -297,64 +297,67 @@ CoulombP3M::CoulombP3M(P3MParameters &&parameters, double prefactor,
set_prefactor(prefactor);
}

namespace {
template <int cao> struct AssignCharge {
void operator()(p3m_data_struct &p3m, double q,
Utils::Vector3d const &real_pos,
p3m_interpolation_cache &inter_weights) {
auto const w = p3m_calculate_interpolation_weights<cao>(
real_pos, p3m.params.ai, p3m.local_mesh);

inter_weights.store(w);

p3m_interpolate(p3m.local_mesh, w, [q, &p3m](int ind, double w) {
p3m.rs_mesh[ind] += w * q;
});
}

void operator()(p3m_data_struct &p3m, double q,
Utils::Vector3d const &real_pos) {
p3m_interpolate(
p3m.local_mesh,
p3m_calculate_interpolation_weights<cao>(real_pos, p3m.params.ai,
p3m.local_mesh),
[q, &p3m](int ind, double w) { p3m.rs_mesh[ind] += w * q; });
}

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 p_q = boost::get<0>(zipped);
auto &p_pos = boost::get<1>(zipped);
if (p_q != 0.0) {
this->operator()(p3m, p_q, p_pos, p3m.inter_weights);
}
}
}
};
} // namespace

template <typename combined_ranges>
void CoulombP3M::charge_assign(combined_ranges const &p_q_pos_range) {
p3m.inter_weights.reset(p3m.params.cao);

/* prepare local FFT mesh */
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,
p_q_pos_range);
}
// TODO: See in hpp file
//
// namespace detail {
// template <int cao> struct AssignCharge {
// void operator()(p3m_data_struct &p3m, double q,
// Utils::Vector3d const &real_pos,
// p3m_interpolation_cache &inter_weights) {
// auto const w = p3m_calculate_interpolation_weights<cao>(
// real_pos, p3m.params.ai, p3m.local_mesh);
//
// inter_weights.store(w);
//
// p3m_interpolate(p3m.local_mesh, w, [q, &p3m](int ind, double w) {
// p3m.rs_mesh[ind] += w * q;
// });
// }
//
// void operator()(p3m_data_struct &p3m, double q,
// Utils::Vector3d const &real_pos) {
// p3m_interpolate(
// p3m.local_mesh,
// p3m_calculate_interpolation_weights<cao>(real_pos, p3m.params.ai,
// p3m.local_mesh),
// [q, &p3m](int ind, double w) { p3m.rs_mesh[ind] += w * q; });
// }
//
// 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 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);
// }
// }
// }
// };
// } // namespace detail

// template <typename combined_ranges>
// void CoulombP3M::charge_assign(combined_ranges const &p_q_pos_range) {
// p3m.inter_weights.reset(p3m.params.cao);
//
// /* prepare local FFT mesh */
// for (int i = 0; i < p3m.local_mesh.size; i++)
// p3m.rs_mesh[i] = 0.0;
//
// Utils::integral_parameter<int, detail::AssignCharge, 1, 7>(
// p3m.params.cao, p3m, p_q_pos_range);
// }

void CoulombP3M::assign_charge(double q, Utils::Vector3d const &real_pos,
p3m_interpolation_cache &inter_weights) {
Utils::integral_parameter<int, AssignCharge, 1, 7>(p3m.params.cao, p3m, q,
real_pos, inter_weights);
Utils::integral_parameter<int, detail::AssignCharge, 1, 7>(
p3m.params.cao, p3m, q, real_pos, inter_weights);
}

void CoulombP3M::assign_charge(double q, Utils::Vector3d const &real_pos) {
Utils::integral_parameter<int, AssignCharge, 1, 7>(p3m.params.cao, p3m, q,
real_pos);
Utils::integral_parameter<int, detail::AssignCharge, 1, 7>(p3m.params.cao,
p3m, q, real_pos);
}

template <int cao> struct AssignForces {
Expand Down Expand Up @@ -475,6 +478,8 @@ double CoulombP3M::long_range_kernel(bool force_flag, bool energy_flag,
p3m.sm.gather_grid(p3m.rs_mesh.data(), comm_cart, p3m.local_mesh.dim);
fft_perform_forw(p3m.rs_mesh.data(), p3m.fft, comm_cart);

// TODO: Move construction up
//
auto p_q_range = ParticlePropertyRange::charge_range(particles);
auto p_force_range = ParticlePropertyRange::force_range(particles);
auto p_unfolded_pos_range =
Expand Down
54 changes: 53 additions & 1 deletion src/core/electrostatics/p3m.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,15 +45,18 @@
#include "p3m/fft.hpp"
#include "p3m/interpolation.hpp"
#include "p3m/send_mesh.hpp"
#include <utils/integral_parameter.hpp>

#include "ParticleRange.hpp"
#include "boost/tuple/tuple.hpp"

#include <utils/Vector.hpp>
#include <utils/constants.hpp>
#include <utils/math/AS_erfc_part.hpp>

#include <array>
#include <cmath>
#include <optional>

struct p3m_data_struct : public p3m_data_struct_base {
explicit p3m_data_struct(P3MParameters &&parameters)
Expand Down Expand Up @@ -81,6 +84,46 @@ struct p3m_data_struct : public p3m_data_struct_base {
fft_data_struct fft;
};

// TODO: Discuss with JN whether we should think about how to move this to the
// cpp file
//
namespace detail {
template <int cao> struct AssignCharge {
void operator()(p3m_data_struct &p3m, double q,
Utils::Vector3d const &real_pos,
p3m_interpolation_cache &inter_weights) {
auto const w = p3m_calculate_interpolation_weights<cao>(
real_pos, p3m.params.ai, p3m.local_mesh);

inter_weights.store(w);

p3m_interpolate(p3m.local_mesh, w, [q, &p3m](int ind, double w) {
p3m.rs_mesh[ind] += w * q;
});
}

void operator()(p3m_data_struct &p3m, double q,
Utils::Vector3d const &real_pos) {
p3m_interpolate(
p3m.local_mesh,
p3m_calculate_interpolation_weights<cao>(real_pos, p3m.params.ai,
p3m.local_mesh),
[q, &p3m](int ind, double w) { p3m.rs_mesh[ind] += w * q; });
}

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);
}
}
}
};
} // namespace detail

/** @brief P3M solver. */
struct CoulombP3M : public Coulomb::Actor<CoulombP3M> {
/** P3M parameters. */
Expand Down Expand Up @@ -164,7 +207,16 @@ struct CoulombP3M : public Coulomb::Actor<CoulombP3M> {
* function.
*/
template <typename combined_ranges>
void charge_assign(combined_ranges const &p_q_pos_range);
void charge_assign(combined_ranges const &p_q_pos_range) {
p3m.inter_weights.reset(p3m.params.cao);

/* prepare local FFT mesh */
for (int i = 0; i < p3m.local_mesh.size; i++)
p3m.rs_mesh[i] = 0.0;

Utils::integral_parameter<int, detail::AssignCharge, 1, 7>(
p3m.params.cao, p3m, p_q_pos_range);
}

/**
* @brief Assign a single charge into the current charge grid.
Expand Down

0 comments on commit 36b27fd

Please sign in to comment.