From: Berk Hess <hess@kth.se>
Date: Fri, 31 Jul 2020 08:47:03 +0000 (+0000)
Subject: Move foreign potential energy accumulation
X-Git-Url: http://biod.pnpi.spb.ru/gitweb/?a=commitdiff_plain;h=99be20fac95e5ba3f4a002bc0b77500ad3e30b02;p=alexxy%2Fgromacs.git

Move foreign potential energy accumulation

The potential energy contributions to the foreign lambda
Hamiltonian differences were summed in sum_dhdl() which meant
that the foreign energy differences were not complete when
do_force() returns. Now there are summed in
accumulatePotentialEnergies() which is called instead of sum_epot().

Also consistently changed the lambda vector to ArrayRef in the force
routines and made it const.
---

diff --git a/src/gromacs/listed_forces/listed_forces.cpp b/src/gromacs/listed_forces/listed_forces.cpp
index e88f050fbd..c35a06a895 100644
--- a/src/gromacs/listed_forces/listed_forces.cpp
+++ b/src/gromacs/listed_forces/listed_forces.cpp
@@ -50,6 +50,7 @@
 
 #include <algorithm>
 #include <array>
+#include <numeric>
 
 #include "gromacs/gmxlib/network.h"
 #include "gromacs/gmxlib/nrnb.h"
@@ -706,7 +707,7 @@ void ListedForces::calculate(struct gmx_wallcycle*          wcycle,
             {
                 gmx_incons("The bonded interactions are not sorted for free energy");
             }
-            for (size_t i = 0; i < enerd->enerpart_lambda.size(); i++)
+            for (int i = 0; i < 1 + enerd->foreignLambdaTerms.numLambdas(); i++)
             {
                 real lam_i[efptNR];
 
@@ -719,12 +720,9 @@ void ListedForces::calculate(struct gmx_wallcycle*          wcycle,
                                    shiftForceBufferLambda_, &(enerd->foreign_grpp), enerd->foreign_term,
                                    dvdl, nrnb, lam_i, md, &fcdata, global_atom_index);
                 sum_epot(enerd->foreign_grpp, enerd->foreign_term);
-                enerd->enerpart_lambda[i] += enerd->foreign_term[F_EPOT];
-                for (int j = 0; j < efptNR; j++)
-                {
-                    enerd->dhdlLambda[i] += dvdl[j];
-                    dvdl[j] = 0;
-                }
+                const double dvdlSum = std::accumulate(std::begin(dvdl), std::end(dvdl), 0.);
+                std::fill(std::begin(dvdl), std::end(dvdl), 0.0);
+                enerd->foreignLambdaTerms.accumulate(i, enerd->foreign_term[F_EPOT], dvdlSum);
             }
             wallcycle_sub_stop(wcycle, ewcsLISTED_FEP);
         }
diff --git a/src/gromacs/listed_forces/position_restraints.cpp b/src/gromacs/listed_forces/position_restraints.cpp
index 40342a9ea7..5b12681b79 100644
--- a/src/gromacs/listed_forces/position_restraints.cpp
+++ b/src/gromacs/listed_forces/position_restraints.cpp
@@ -434,21 +434,20 @@ void posres_wrapper_lambda(struct gmx_wallcycle*         wcycle,
                            const real*                   lambda,
                            const t_forcerec*             fr)
 {
-    real v;
-
     wallcycle_sub_start_nocount(wcycle, ewcsRESTRAINTS);
-    for (size_t i = 0; i < enerd->enerpart_lambda.size(); i++)
+
+    auto& foreignTerms = enerd->foreignLambdaTerms;
+    for (int i = 0; i < 1 + foreignTerms.numLambdas(); i++)
     {
         real dvdl = 0;
 
         const real lambda_dum =
                 (i == 0 ? lambda[efptRESTRAINT] : fepvals->all_lambda[efptRESTRAINT][i - 1]);
-        v = posres<false>(idef.il[F_POSRES].size(), idef.il[F_POSRES].iatoms.data(),
-                          idef.iparams_posres.data(), x, nullptr,
-                          fr->pbcType == PbcType::No ? nullptr : pbc, lambda_dum, &dvdl,
-                          fr->rc_scaling, fr->pbcType, fr->posres_com, fr->posres_comB);
-        enerd->enerpart_lambda[i] += v;
-        enerd->dhdlLambda[i] += dvdl;
+        const real v = posres<false>(idef.il[F_POSRES].size(), idef.il[F_POSRES].iatoms.data(),
+                                     idef.iparams_posres.data(), x, nullptr,
+                                     fr->pbcType == PbcType::No ? nullptr : pbc, lambda_dum, &dvdl,
+                                     fr->rc_scaling, fr->pbcType, fr->posres_com, fr->posres_comB);
+        foreignTerms.accumulate(i, v, dvdl);
     }
     wallcycle_sub_stop(wcycle, ewcsRESTRAINTS);
 }
diff --git a/src/gromacs/mdlib/enerdata_utils.cpp b/src/gromacs/mdlib/enerdata_utils.cpp
index 4d9a758bb1..fae35d5e0d 100644
--- a/src/gromacs/mdlib/enerdata_utils.cpp
+++ b/src/gromacs/mdlib/enerdata_utils.cpp
@@ -39,15 +39,46 @@
 
 #include "enerdata_utils.h"
 
+#include "gromacs/gmxlib/network.h"
+#include "gromacs/mdtypes/commrec.h"
 #include "gromacs/mdtypes/enerdata.h"
 #include "gromacs/mdtypes/inputrec.h"
 #include "gromacs/utility/fatalerror.h"
 #include "gromacs/utility/smalloc.h"
 
+ForeignLambdaTerms::ForeignLambdaTerms(int numLambdas) :
+    numLambdas_(numLambdas),
+    energies_(1 + numLambdas),
+    dhdl_(1 + numLambdas)
+{
+}
+
+std::pair<std::vector<double>, std::vector<double>> ForeignLambdaTerms::getTerms(t_commrec* cr)
+{
+    std::vector<double> data(2 * numLambdas_);
+    for (int i = 0; i < numLambdas_; i++)
+    {
+        data[i]               = energies_[1 + i] - energies_[0];
+        data[numLambdas_ + i] = dhdl_[1 + i];
+    }
+    if (cr && cr->nnodes > 1)
+    {
+        gmx_sumd(data.size(), data.data(), cr);
+    }
+    auto dataMid = data.begin() + numLambdas_;
+
+    return { { data.begin(), dataMid }, { dataMid, data.end() } };
+}
+
+void ForeignLambdaTerms::zeroAllTerms()
+{
+    std::fill(energies_.begin(), energies_.end(), 0.0);
+    std::fill(dhdl_.begin(), dhdl_.end(), 0.0);
+}
+
 gmx_enerdata_t::gmx_enerdata_t(int numEnergyGroups, int numFepLambdas) :
     grpp(numEnergyGroups),
-    enerpart_lambda(numFepLambdas == 0 ? 0 : numFepLambdas + 1),
-    dhdlLambda(numFepLambdas == 0 ? 0 : numFepLambdas + 1),
+    foreignLambdaTerms(numFepLambdas),
     foreign_grpp(numEnergyGroups)
 {
 }
@@ -91,27 +122,17 @@ void sum_epot(const gmx_grppairener_t& grpp, real* epot)
     }
 }
 
-/* Adds computed dV/dlambda contributions to the requested outputs
- *
- * Also adds the dispersion correction dV/dlambda to the VdW term.
- * Note that kinetic energy and constraint contributions are handled in sum_dkdl()
- */
-static void sum_dvdl(gmx_enerdata_t* enerd, const t_lambda& fepvals)
+// Adds computed dV/dlambda contributions to the requested outputs
+static void set_dvdl_output(gmx_enerdata_t* enerd, const t_lambda& fepvals)
 {
-    // Add dispersion correction to the VdW component
-    enerd->dvdl_lin[efptVDW] += enerd->term[F_DVDL_VDW];
-
-    for (size_t i = 0; i < enerd->enerpart_lambda.size(); i++)
-    {
-        enerd->dhdlLambda[i] += enerd->term[F_DVDL_VDW];
-    }
-
     enerd->term[F_DVDL] = 0.0;
     for (int i = 0; i < efptNR; i++)
     {
         if (fepvals.separate_dvdl[i])
         {
-            /* could this be done more readably/compactly? */
+            /* Translate free-energy term indices to idef term indices.
+             * Could this be done more readably/compactly?
+             */
             int index;
             switch (i)
             {
@@ -145,39 +166,47 @@ void accumulatePotentialEnergies(gmx_enerdata_t* enerd, gmx::ArrayRef<const real
 {
     sum_epot(enerd->grpp, enerd->term);
 
-    if (fepvals)
+    if (fepvals == nullptr)
     {
-        /* Note that sum_dvdl() adds the dispersion correction enerd->dvdl_lin[efptVDW],
-         * so sum_dvdl() should be called before computing the foreign lambda energy differences.
-         */
-        sum_dvdl(enerd, *fepvals);
+        return;
+    }
 
-        /* Sum the foreign lambda energy difference contributions.
-         * Note that here we only add the potential energy components.
-         * The constraint and kinetic energy components are add after integration
-         * by sum_dhdl().
-         */
-        for (int i = 0; i < fepvals->n_lambda; i++)
-        {
-            /* note we are iterating over fepvals here!
-               For the current lam, dlam = 0 automatically,
-               so we don't need to add anything to the
-               enerd->enerpart_lambda[0] */
+    // Accumulate the foreign lambda terms
 
-            /* we don't need to worry about dvdl_lin contributions to dE at
-               current lambda, because the contributions to the current
-               lambda are automatically zeroed */
+    double dvdl_lin = 0;
+    for (int i = 0; i < efptNR; i++)
+    {
+        dvdl_lin += enerd->dvdl_lin[i];
+    }
+    enerd->foreignLambdaTerms.addConstantDhdl(dvdl_lin);
 
-            double& enerpart_lambda = enerd->enerpart_lambda[i + 1];
+    set_dvdl_output(enerd, *fepvals);
 
-            for (gmx::index j = 0; j < lambda.ssize(); j++)
-            {
-                /* Note that this loop is over all dhdl components, not just the separated ones */
-                const double dlam = fepvals->all_lambda[j][i] - lambda[j];
+    /* Sum the foreign lambda energy difference contributions.
+     * Note that here we only add the potential energy components.
+     * The constraint and kinetic energy components are add after integration
+     * by sum_dhdl().
+     */
+    for (int i = 0; i < fepvals->n_lambda; i++)
+    {
+        /* note we are iterating over fepvals here!
+           For the current lam, dlam = 0 automatically,
+           so we don't need to add anything to the
+           enerd->enerpart_lambda[0] */
 
-                enerpart_lambda += dlam * enerd->dvdl_lin[j];
-            }
+        /* we don't need to worry about dvdl_lin contributions to dE at
+           current lambda, because the contributions to the current
+           lambda are automatically zeroed */
+
+        double enerpart_lambda = 0;
+        for (gmx::index j = 0; j < lambda.ssize(); j++)
+        {
+            /* Note that this loop is over all dhdl components, not just the separated ones */
+            const double dlam = fepvals->all_lambda[j][i] - lambda[j];
+
+            enerpart_lambda += dlam * enerd->dvdl_lin[j];
         }
+        enerd->foreignLambdaTerms.accumulate(1 + i, enerpart_lambda, 0);
     }
 }
 
@@ -194,15 +223,15 @@ void accumulateKineticLambdaComponents(gmx_enerdata_t*           enerd,
         enerd->term[F_DVDL] += enerd->term[F_DVDL_CONSTR];
     }
 
-    for (int i = 0; i < fepvals.n_lambda; i++)
+    // Treat current lambda, which only needs a dV/dl contribution
+    enerd->foreignLambdaTerms.accumulate(0, 0.0, enerd->term[F_DVDL_CONSTR]);
+    if (!fepvals.separate_dvdl[efptMASS])
     {
-        /* note we are iterating over fepvals here!
-           For the current lam, dlam = 0 automatically,
-           so we don't need to add anything to the
-           enerd->enerpart_lambda[0] */
-
-        double& enerpart_lambda = enerd->enerpart_lambda[i + 1];
+        enerd->foreignLambdaTerms.accumulate(0, 0.0, enerd->term[F_DKDL]);
+    }
 
+    for (int i = 0; i < fepvals.n_lambda; i++)
+    {
         /* Note that potential energy terms have been added by sum_epot() -> sum_dvdl() */
 
         /* Constraints can not be evaluated at foreign lambdas, so we add
@@ -211,12 +240,13 @@ void accumulateKineticLambdaComponents(gmx_enerdata_t*           enerd,
          */
         const int    lambdaIndex = (fepvals.separate_dvdl[efptBONDED] ? efptBONDED : efptFEP);
         const double dlam        = fepvals.all_lambda[lambdaIndex][i] - lambda[lambdaIndex];
-        enerpart_lambda += dlam * enerd->term[F_DVDL_CONSTR];
+        enerd->foreignLambdaTerms.accumulate(1 + i, dlam * enerd->term[F_DVDL_CONSTR],
+                                             enerd->term[F_DVDL_CONSTR]);
 
         if (!fepvals.separate_dvdl[efptMASS])
         {
             const double dlam = fepvals.all_lambda[efptMASS][i] - lambda[efptMASS];
-            enerpart_lambda += dlam * enerd->term[F_DKDL];
+            enerd->foreignLambdaTerms.accumulate(1 + i, dlam * enerd->term[F_DKDL], enerd->term[F_DKDL]);
         }
     }
 
@@ -280,8 +310,7 @@ void reset_enerdata(gmx_enerdata_t* enerd)
     enerd->term[F_DVDL_BONDED]    = 0.0_real;
     enerd->term[F_DVDL_RESTRAINT] = 0.0_real;
     enerd->term[F_DKDL]           = 0.0_real;
-    std::fill(enerd->enerpart_lambda.begin(), enerd->enerpart_lambda.end(), 0);
-    std::fill(enerd->dhdlLambda.begin(), enerd->dhdlLambda.end(), 0);
+    enerd->foreignLambdaTerms.zeroAllTerms();
     /* reset foreign energy data and dvdl - separate functions since they are also called elsewhere */
     reset_foreign_enerdata(enerd);
     reset_dvdl_enerdata(enerd);
diff --git a/src/gromacs/mdlib/energyoutput.cpp b/src/gromacs/mdlib/energyoutput.cpp
index 589d88fce3..120591f0a7 100644
--- a/src/gromacs/mdlib/energyoutput.cpp
+++ b/src/gromacs/mdlib/energyoutput.cpp
@@ -1055,16 +1055,17 @@ void EnergyOutput::addDataAtEnergyStep(bool                    bDoDHDL,
     // BAR + thermodynamic integration values
     if ((fp_dhdl_ || dhc_) && bDoDHDL)
     {
-        for (gmx::index i = 0; i < static_cast<gmx::index>(enerd->enerpart_lambda.size()) - 1; i++)
+        const auto& foreignTerms = enerd->foreignLambdaTerms;
+        for (int i = 0; i < foreignTerms.numLambdas(); i++)
         {
             /* zero for simulated tempering */
-            dE_[i] = enerd->enerpart_lambda[i + 1] - enerd->enerpart_lambda[0];
+            dE_[i] = foreignTerms.deltaH(i);
             if (numTemperatures_ > 0)
             {
                 GMX_RELEASE_ASSERT(numTemperatures_ > state->fep_state,
                                    "Number of lambdas in state is bigger then in input record");
                 GMX_RELEASE_ASSERT(
-                        numTemperatures_ >= static_cast<gmx::index>(enerd->enerpart_lambda.size()) - 1,
+                        numTemperatures_ >= foreignTerms.numLambdas(),
                         "Number of lambdas in energy data is bigger then in input record");
                 /* MRS: is this right, given the way we have defined the exchange probabilities? */
                 /* is this even useful to have at all? */
@@ -1111,7 +1112,7 @@ void EnergyOutput::addDataAtEnergyStep(bool                    bDoDHDL,
             {
                 fprintf(fp_dhdl_, " %#.8g", dE_[i]);
             }
-            if (bDynBox_ && bDiagPres_ && (epc_ != epcNO) && !enerd->enerpart_lambda.empty()
+            if (bDynBox_ && bDiagPres_ && (epc_ != epcNO) && foreignTerms.numLambdas() > 0
                 && (fep->init_lambda < 0))
             {
                 fprintf(fp_dhdl_, " %#.8g", pv); /* PV term only needed when
diff --git a/src/gromacs/mdlib/expanded.cpp b/src/gromacs/mdlib/expanded.cpp
index 2d82013c81..9e55df33da 100644
--- a/src/gromacs/mdlib/expanded.cpp
+++ b/src/gromacs/mdlib/expanded.cpp
@@ -2,7 +2,7 @@
  * This file is part of the GROMACS molecular simulation package.
  *
  * Copyright (c) 2012-2018, The GROMACS development team.
- * Copyright (c) 2019, by the GROMACS development team, led by
+ * Copyright (c) 2019,2020, by the GROMACS development team, led by
  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
  * and including many others, as listed in the AUTHORS file in the
  * top-level source directory and at http://www.gromacs.org.
@@ -1248,8 +1248,7 @@ int ExpandedEnsembleDynamics(FILE*                 log,
             if (ir->bSimTemp)
             {
                 /* Note -- this assumes no mass changes, since kinetic energy is not added  . . . */
-                scaled_lamee[i] = (enerd->enerpart_lambda[i + 1] - enerd->enerpart_lambda[0])
-                                          / (simtemp->temperatures[i] * BOLTZ)
+                scaled_lamee[i] = enerd->foreignLambdaTerms.deltaH(i) / (simtemp->temperatures[i] * BOLTZ)
                                   + enerd->term[F_EPOT]
                                             * (1.0 / (simtemp->temperatures[i])
                                                - 1.0 / (simtemp->temperatures[fep_state]))
@@ -1257,8 +1256,7 @@ int ExpandedEnsembleDynamics(FILE*                 log,
             }
             else
             {
-                scaled_lamee[i] = (enerd->enerpart_lambda[i + 1] - enerd->enerpart_lambda[0])
-                                  / (expand->mc_temp * BOLTZ);
+                scaled_lamee[i] = enerd->foreignLambdaTerms.deltaH(i) / (expand->mc_temp * BOLTZ);
                 /* mc_temp is currently set to the system reft unless otherwise defined */
             }
 
diff --git a/src/gromacs/mdlib/force.cpp b/src/gromacs/mdlib/force.cpp
index dd5887ced2..6805a79454 100644
--- a/src/gromacs/mdlib/force.cpp
+++ b/src/gromacs/mdlib/force.cpp
@@ -132,11 +132,6 @@ void do_force_lowlevel(t_forcerec*                          fr,
         real dvdl_walls = do_walls(*ir, *fr, box, *md, x, &forceWithVirial, lambda[efptVDW],
                                    enerd->grpp.ener[egLJSR].data(), nrnb);
         enerd->dvdl_lin[efptVDW] += dvdl_walls;
-
-        for (auto& dhdl : enerd->dhdlLambda)
-        {
-            dhdl += dvdl_walls;
-        }
     }
 
     {
@@ -294,11 +289,6 @@ void do_force_lowlevel(t_forcerec*                          fr,
         enerd->term[F_COUL_RECIP] = Vlr_q + ewaldOutput.Vcorr_q;
         enerd->term[F_LJ_RECIP]   = Vlr_lj + ewaldOutput.Vcorr_lj;
 
-        for (auto& dhdl : enerd->dhdlLambda)
-        {
-            dhdl += ewaldOutput.dvdl[efptVDW] + ewaldOutput.dvdl[efptCOUL];
-        }
-
         if (debug)
         {
             fprintf(debug, "Vlr_q = %g, Vcorr_q = %g, Vlr_corr_q = %g\n", Vlr_q,
diff --git a/src/gromacs/mdlib/md_support.cpp b/src/gromacs/mdlib/md_support.cpp
index 9b04614c8b..d6b7947435 100644
--- a/src/gromacs/mdlib/md_support.cpp
+++ b/src/gromacs/mdlib/md_support.cpp
@@ -484,10 +484,7 @@ void compute_globals(gmx_global_stat*               gstat,
 
         enerd->term[F_EKIN] = trace(ekind->ekin);
 
-        for (auto& dhdl : enerd->dhdlLambda)
-        {
-            dhdl += enerd->dvdl_lin[efptMASS];
-        }
+        enerd->foreignLambdaTerms.addConstantDhdl(enerd->dvdl_lin[efptMASS]);
     }
 
     /* ########## Now pressure ############## */
diff --git a/src/gromacs/mdlib/sim_util.cpp b/src/gromacs/mdlib/sim_util.cpp
index 611ca913dc..1d61a68c59 100644
--- a/src/gromacs/mdlib/sim_util.cpp
+++ b/src/gromacs/mdlib/sim_util.cpp
@@ -201,10 +201,6 @@ static void pull_potential_wrapper(const t_commrec*               cr,
             pull_potential(pull_work, mdatoms->massT, &pbc, cr, t, lambda[efptRESTRAINT],
                            as_rvec_array(x.data()), force, &dvdl);
     enerd->dvdl_lin[efptRESTRAINT] += dvdl;
-    for (auto& dhdl : enerd->dhdlLambda)
-    {
-        dhdl += dvdl;
-    }
     wallcycle_stop(wcycle, ewcPULLPOT);
 }
 
@@ -235,11 +231,6 @@ static void pme_receive_force_ener(t_forcerec*           fr,
     enerd->dvdl_lin[efptCOUL] += dvdl_q;
     enerd->dvdl_lin[efptVDW] += dvdl_lj;
 
-    for (auto& dhdl : enerd->dhdlLambda)
-    {
-        dhdl += dvdl_q + dvdl_lj;
-    }
-
     if (wcycle)
     {
         dd_cycles_add(cr->dd, cycles_seppme, ddCyclPME);
@@ -1865,6 +1856,7 @@ void do_force(FILE*                               fplog,
         {
             enerd->term[F_DISPCORR] = correction.energy;
             enerd->term[F_DVDL_VDW] += correction.dvdl;
+            enerd->dvdl_lin[efptVDW] += correction.dvdl;
         }
         if (stepWork.computeVirial)
         {
@@ -1894,7 +1886,6 @@ void do_force(FILE*                               fplog,
     {
         /* Compute the final potential energy terms */
         accumulatePotentialEnergies(enerd, lambda, inputrec->fepvals);
-        ;
 
         if (!EI_TPI(inputrec->eI))
         {
diff --git a/src/gromacs/mdlib/stat.cpp b/src/gromacs/mdlib/stat.cpp
index 99619beab0..60d6445f2d 100644
--- a/src/gromacs/mdlib/stat.cpp
+++ b/src/gromacs/mdlib/stat.cpp
@@ -256,9 +256,10 @@ void global_stat(const gmx_global_stat*  gs,
         {
             idvdll  = add_bind(rb, efptNR, enerd->dvdl_lin);
             idvdlnl = add_bind(rb, efptNR, enerd->dvdl_nonlin);
-            if (!enerd->enerpart_lambda.empty())
+            if (enerd->foreignLambdaTerms.numLambdas() > 0)
             {
-                iepl = add_bind(rb, enerd->enerpart_lambda.size(), enerd->enerpart_lambda.data());
+                iepl = add_bind(rb, enerd->foreignLambdaTerms.energies().size(),
+                                enerd->foreignLambdaTerms.energies().data());
             }
         }
     }
@@ -351,9 +352,10 @@ void global_stat(const gmx_global_stat*  gs,
         {
             extract_bind(rb, idvdll, efptNR, enerd->dvdl_lin);
             extract_bind(rb, idvdlnl, efptNR, enerd->dvdl_nonlin);
-            if (!enerd->enerpart_lambda.empty())
+            if (enerd->foreignLambdaTerms.numLambdas() > 0)
             {
-                extract_bind(rb, iepl, enerd->enerpart_lambda.size(), enerd->enerpart_lambda.data());
+                extract_bind(rb, iepl, enerd->foreignLambdaTerms.energies().size(),
+                             enerd->foreignLambdaTerms.energies().data());
             }
         }
 
diff --git a/src/gromacs/mdrun/replicaexchange.cpp b/src/gromacs/mdrun/replicaexchange.cpp
index 4c036dc2fb..c40161d9ef 100644
--- a/src/gromacs/mdrun/replicaexchange.cpp
+++ b/src/gromacs/mdrun/replicaexchange.cpp
@@ -919,8 +919,7 @@ static void test_for_replica_exchange(FILE*                 fplog,
         }
         for (i = 0; i < re->nrepl; i++)
         {
-            re->de[i][re->repl] = (enerd->enerpart_lambda[static_cast<int>(re->q[ereLAMBDA][i]) + 1]
-                                   - enerd->enerpart_lambda[0]);
+            re->de[i][re->repl] = enerd->foreignLambdaTerms.deltaH(re->q[ereLAMBDA][i]);
         }
     }
 
diff --git a/src/gromacs/mdtypes/enerdata.h b/src/gromacs/mdtypes/enerdata.h
index 9866d7740b..12c3a051ae 100644
--- a/src/gromacs/mdtypes/enerdata.h
+++ b/src/gromacs/mdtypes/enerdata.h
@@ -36,12 +36,17 @@
 #define GMX_MDTYPES_TYPES_ENERDATA_H
 
 #include <array>
+#include <utility>
 #include <vector>
 
 #include "gromacs/mdtypes/md_enums.h"
 #include "gromacs/topology/idef.h"
+#include "gromacs/utility/arrayref.h"
 #include "gromacs/utility/real.h"
 
+struct t_commrec;
+
+// The non-bonded energy terms accumulated for energy group pairs
 enum
 {
     egCOULSR,
@@ -52,6 +57,7 @@ enum
     egNR
 };
 
+// Struct for accumulating non-bonded energies between energy group pairs
 struct gmx_grppairener_t
 {
     gmx_grppairener_t(int numEnergyGroups) : nener(numEnergyGroups * numEnergyGroups)
@@ -66,24 +72,103 @@ struct gmx_grppairener_t
     std::array<std::vector<real>, egNR> ener;  /* Energy terms for each pair of groups */
 };
 
+//! Accumulates free-energy foreign lambda energies and dH/dlamba
+class ForeignLambdaTerms
+{
+public:
+    /*! \brief Constructor
+     *
+     * \param[in] numLambdas  The number of foreign lambda values
+     */
+    ForeignLambdaTerms(int numLambdas);
+
+    //! Returns the number of foreign lambda values
+    int numLambdas() const { return numLambdas_; }
+
+    //! Returns the H(lambdaIndex) - H(lambda_current)
+    double deltaH(int lambdaIndex) const { return energies_[1 + lambdaIndex] - energies_[0]; }
+
+    /*! \brief Returns a list of partial energies, the part which depends on lambda),
+     * current lambda in entry 0, foreign lambda i in entry 1+i
+     */
+    gmx::ArrayRef<double> energies() { return energies_; }
+
+    /*! \brief Returns a list of partial energies, the part which depends on lambda),
+     * current lambda in entry 0, foreign lambda i in entry 1+i
+     */
+    gmx::ArrayRef<const double> energies() const { return energies_; }
+
+    /*! \brief Adds an energy and dH/dl constribution to lambda list index \p listIndex
+     *
+     * This should only be used for terms with non-linear dependence on lambda
+     * The value passed as listIndex should be 0 for the current lambda
+     * and 1+i for foreign lambda index i.
+     */
+    void accumulate(int listIndex, double energy, double dhdl)
+    {
+        energies_[listIndex] += energy;
+        dhdl_[listIndex] += dhdl;
+    }
+
+    /*! \brief Add a dH/dl contribution that does not depend on lambda to all foreign dH/dl terms
+     *
+     * Note: this should not be called directly for energy terms that depend linearly on lambda,
+     * as those are added automatically through the accumulated dvdl_lin term in gmx_enerdata_t.
+     */
+    void addConstantDhdl(double dhdl)
+    {
+        for (double& foreignDhdl : dhdl_)
+        {
+            foreignDhdl += dhdl;
+        }
+    }
+
+    /*! \brief Returns a pair of lists of deltaH and dH/dlambda
+     *
+     * Both lists are of size numLambdas() and are indexed with the lambda index.
+     * The returned lists are valid until the next call to this method.
+     *
+     * \param[in] cr  Communication record, used to reduce the terms when !=nullptr
+     */
+    std::pair<std::vector<double>, std::vector<double>> getTerms(t_commrec* cr);
+
+    //! Sets all terms to 0
+    void zeroAllTerms();
+
+private:
+    //! The number of foreign lambdas
+    int numLambdas_;
+    //! Storage for foreign lambda energies
+    std::vector<double> energies_;
+    //! Storage for foreign lambda dH/dlambda
+    std::vector<double> dhdl_;
+};
+
+//! Struct for accumulating all potential energy terms and some kinetic energy terms
 struct gmx_enerdata_t
 {
     gmx_enerdata_t(int numEnergyGroups, int numFepLambdas);
 
-    real term[F_NRE] = { 0 }; /* The energies for all different interaction types */
+    //! The energies for all different interaction types
+    real term[F_NRE] = { 0 };
+    //! Energy group pair non-bonded energies
     struct gmx_grppairener_t grpp;
-    double dvdl_lin[efptNR]    = { 0 }; /* Contributions to dvdl with linear lam-dependence */
-    double dvdl_nonlin[efptNR] = { 0 }; /* Idem, but non-linear dependence                  */
+    //! Contributions to dV/dlambda with linear dependence on lambda
+    double dvdl_lin[efptNR] = { 0 };
+    //! Contributions to dV/dlambda with non-linear dependence on lambda
+    double dvdl_nonlin[efptNR] = { 0 };
     /* The idea is that dvdl terms with linear lambda dependence will be added
      * automatically to enerpart_lambda. Terms with non-linear lambda dependence
      * should explicitly determine the energies at foreign lambda points
      * when n_lambda > 0. */
 
-    int                 fep_state = 0; /*current fep state -- just for printing */
-    std::vector<double> enerpart_lambda; /* Partial Hamiltonian for lambda and flambda[], includes at least all perturbed terms */
-    std::vector<double> dhdlLambda; /* dHdL at all neighboring lambda points (the current lambda point also at index 0). */
-    real foreign_term[F_NRE] = { 0 };      /* alternate array for storing foreign lambda energies */
-    struct gmx_grppairener_t foreign_grpp; /* alternate array for storing foreign lambda energies */
+    //! Foreign lambda energies and dH/dl
+    ForeignLambdaTerms foreignLambdaTerms;
+
+    //! Alternate, temporary array for storing foreign lambda energies
+    real foreign_term[F_NRE] = { 0 };
+    //! Alternate, temporary  array for storing foreign lambda group pair energies
+    struct gmx_grppairener_t foreign_grpp;
 };
 
 #endif
diff --git a/src/gromacs/nbnxm/kerneldispatch.cpp b/src/gromacs/nbnxm/kerneldispatch.cpp
index ea47bb2090..0eeeebc68a 100644
--- a/src/gromacs/nbnxm/kerneldispatch.cpp
+++ b/src/gromacs/nbnxm/kerneldispatch.cpp
@@ -538,7 +538,7 @@ void nonbonded_verlet_t::dispatchFreeEnergyKernel(gmx::InteractionLocality   iLo
         kernel_data.energygrp_elec = enerd->foreign_grpp.ener[egCOULSR].data();
         kernel_data.energygrp_vdw  = enerd->foreign_grpp.ener[egLJSR].data();
 
-        for (size_t i = 0; i < enerd->enerpart_lambda.size(); i++)
+        for (gmx::index i = 0; i < enerd->foreignLambdaTerms.energies().ssize(); i++)
         {
             std::fill(std::begin(dvdl_nb), std::end(dvdl_nb), 0);
             for (int j = 0; j < efptNR; j++)
@@ -558,15 +558,8 @@ void nonbonded_verlet_t::dispatchFreeEnergyKernel(gmx::InteractionLocality   iLo
             }
 
             sum_epot(enerd->foreign_grpp, enerd->foreign_term);
-            enerd->enerpart_lambda[i] += enerd->foreign_term[F_EPOT];
-            enerd->dhdlLambda[i] += dvdl_nb[efptVDW] + dvdl_nb[efptCOUL];
-        }
-    }
-    else
-    {
-        for (size_t i = 0; i < enerd->enerpart_lambda.size(); i++)
-        {
-            enerd->dhdlLambda[i] += dvdl_nb[efptVDW] + dvdl_nb[efptCOUL];
+            enerd->foreignLambdaTerms.accumulate(i, enerd->foreign_term[F_EPOT],
+                                                 dvdl_nb[efptVDW] + dvdl_nb[efptCOUL]);
         }
     }
     wallcycle_sub_stop(wcycle_, ewcsNONBONDED_FEP);