diff --git a/src/engine.c b/src/engine.c
index a323633636a6b276c7e16409d58c6ffb457a23f1..33b49c5ee49af77862ebe275a043c2ca3dd3279b 100644
--- a/src/engine.c
+++ b/src/engine.c
@@ -6075,6 +6075,16 @@ void engine_config(int restart, struct engine *e, struct swift_params *params,
e->time_first_statistics, e->time_begin);
}
+ /* Get the total mass */
+ e->total_mass = 0.;
+ for (size_t i = 0; i < e->s->nr_gparts; ++i)
+ e->total_mass += e->s->gparts[i].mass;
+
+/* Reduce the total mass */
+#ifdef WITH_MPI
+ MPI_Allreduce(MPI_IN_PLACE, &e->total_mass, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
+#endif
+
/* Find the time of the first snapshot output */
engine_compute_next_snapshot_time(e);
diff --git a/src/engine.h b/src/engine.h
index b29b5bb0bb6013e14c49f0d45bd072984822e012..d5de93c8c1950321581b0961db7039f158083859 100644
--- a/src/engine.h
+++ b/src/engine.h
@@ -200,6 +200,9 @@ struct engine {
/* Total numbers of particles in the system. */
long long total_nr_parts, total_nr_gparts, total_nr_sparts;
+ /* Total mass in the simulation */
+ double total_mass;
+
/* The internal system of units */
const struct unit_system *internal_units;
diff --git a/src/gravity.c b/src/gravity.c
index d0eb589d15ba89aecc5a581de0c143dc7504d3ee..24aa7f56b9a950227b04a92dde1ec5d9af4a84ed 100644
--- a/src/gravity.c
+++ b/src/gravity.c
@@ -152,6 +152,7 @@ void gravity_exact_force_ewald_init(double boxSize) {
bzero(fewald_z, (Newald + 1) * (Newald + 1) * (Newald + 1) * sizeof(float));
bzero(potewald, (Newald + 1) * (Newald + 1) * (Newald + 1) * sizeof(float));
+ /* Hernquist, Bouchet & Suto, 1991, Eq. 2.10 and just below Eq. 2.15 */
potewald[0][0][0] = 2.8372975f;
/* Compute the values in one of the octants */
diff --git a/src/gravity/Default/gravity.h b/src/gravity/Default/gravity.h
index c90ca39e53c525635fe328312babf84f266dfd9a..7e8139ea93ea16eb05c5ede75291021295c78b80 100644
--- a/src/gravity/Default/gravity.h
+++ b/src/gravity/Default/gravity.h
@@ -161,13 +161,16 @@ __attribute__((always_inline)) INLINE static void gravity_init_gpart(
/**
* @brief Finishes the gravity calculation.
*
- * Multiplies the forces and accelerations by the appropiate constants
+ * Multiplies the forces and accelerations by the appropiate constants.
+ * Applies cosmological correction for periodic BCs.
*
* @param gp The particle to act upon
- * @param const_G Newton's constant in internal units
+ * @param const_G Newton's constant in internal units.
+ * @param potential_normalisation Term to be added to all the particles.
*/
__attribute__((always_inline)) INLINE static void gravity_end_force(
- struct gpart* gp, float const_G) {
+ struct gpart* gp, float const_G, const float potential_normalisation,
+ const int periodic) {
/* Let's get physical... */
gp->a_grav[0] *= const_G;
diff --git a/src/gravity/Potential/gravity.h b/src/gravity/Potential/gravity.h
index 2d1dfb8c8f659e86cda23b4d4bfde216b4f5b28f..494b80f76fc8c21ded2aa3fc0e5229f207174bc7 100644
--- a/src/gravity/Potential/gravity.h
+++ b/src/gravity/Potential/gravity.h
@@ -157,13 +157,20 @@ __attribute__((always_inline)) INLINE static void gravity_init_gpart(
/**
* @brief Finishes the gravity calculation.
*
- * Multiplies the forces and accelerations by the appropiate constants
+ * Multiplies the forces and accelerations by the appropiate constants.
+ * Applies cosmological correction for periodic BCs.
*
* @param gp The particle to act upon
- * @param const_G Newton's constant in internal units
+ * @param const_G Newton's constant in internal units.
+ * @param potential_normalisation Term to be added to all the particles.
*/
__attribute__((always_inline)) INLINE static void gravity_end_force(
- struct gpart* gp, float const_G, float Omega_m, float rho_crit0, const int periodic_cosmology) {
+ struct gpart* gp, float const_G, const float potential_normalisation,
+ const int periodic) {
+
+ /* Apply the periodic correction to the peculiar potential */
+ if(periodic)
+ gp->potential += potential_normalisation;
/* Let's get physical... */
gp->a_grav[0] *= const_G;
@@ -171,12 +178,6 @@ __attribute__((always_inline)) INLINE static void gravity_end_force(
gp->a_grav[2] *= const_G;
gp->potential *= const_G;
- /* Apply the cosmological correction to the potential */
- if(periodic_cosmology) {
- const float mass2 = gp->mass * gp->mass;
- gp->potential -= 2.8372975f * cbrtf(mass2 * Omega_m * rho_crit0);
- }
-
#ifdef SWIFT_GRAVITY_FORCE_CHECKS
gp->potential_PM *= const_G;
gp->a_grav_PM[0] *= const_G;
diff --git a/src/runner.c b/src/runner.c
index 8f0fe5ca4be35835e76de0c7a750574ec22182c5..d4dea7a61b63a0b95ad1fe3c96526066a8f9d773 100644
--- a/src/runner.c
+++ b/src/runner.c
@@ -1663,12 +1663,10 @@ void runner_do_end_force(struct runner *r, struct cell *c, int timer) {
struct gpart *restrict gparts = c->gparts;
struct spart *restrict sparts = c->sparts;
const int periodic = s->periodic;
- const int with_cosmology = e->policy & engine_policy_cosmology;
- const int with_periodic_cosmology = periodic && with_cosmology;
- const float Omega_m = e->cosmology->Omega_m;
- const float H0 = e->cosmology->H0;
const float G_newton = e->physical_constants->const_newton_G;
- const float rho_crit0 = 3.f * H0 * H0 / (8.f * M_PI * G_newton);
+ const double r_s = e->mesh->r_s;
+ const double volume = s->dim[0] * s->dim[1] * s->dim[2];
+ const float potential_normalisation = 4. * M_PI * e->total_mass * r_s * r_s / volume;
TIMER_TIC;
@@ -1703,7 +1701,7 @@ void runner_do_end_force(struct runner *r, struct cell *c, int timer) {
if (gpart_is_active(gp, e)) {
/* Finish the force calculation */
- gravity_end_force(gp, G_newton, Omega_m, rho_crit0, with_periodic_cosmology);
+ gravity_end_force(gp, G_newton, potential_normalisation, periodic);
#ifdef SWIFT_NO_GRAVITY_BELOW_ID