src/gromacs/mdlib/expanded.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2012-2018, The GROMACS development team.
   5  * Copyright (c) 2019,2020,2021, by the GROMACS development team, led by
   6  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   7  * and including many others, as listed in the AUTHORS file in the
   8  * top-level source directory and at http://www.gromacs.org.
   9  *
  10  * GROMACS is free software; you can redistribute it and/or
  11  * modify it under the terms of the GNU Lesser General Public License
  12  * as published by the Free Software Foundation; either version 2.1
  13  * of the License, or (at your option) any later version.
  14  *
  15  * GROMACS is distributed in the hope that it will be useful,
  16  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  17  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  18  * Lesser General Public License for more details.
  19  *
  20  * You should have received a copy of the GNU Lesser General Public
  21  * License along with GROMACS; if not, see
  22  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  23  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  24  *
  25  * If you want to redistribute modifications to GROMACS, please
  26  * consider that scientific software is very special. Version
  27  * control is crucial - bugs must be traceable. We will be happy to
  28  * consider code for inclusion in the official distribution, but
  29  * derived work must not be called official GROMACS. Details are found
  30  * in the README & COPYING files - if they are missing, get the
  31  * official version at http://www.gromacs.org.
  32  *
  33  * To help us fund GROMACS development, we humbly ask that you cite
  34  * the research papers on the package. Check out http://www.gromacs.org.
  35  */
  36 #include "gmxpre.h"
  37
  38 #include "expanded.h"
  39
  40 #include <cmath>
  41 #include <cstdio>
  42
  43 #include "gromacs/math/units.h"
  44 #include "gromacs/math/utilities.h"
  45 #include "gromacs/mdtypes/enerdata.h"
  46 #include "gromacs/mdtypes/forcerec.h"
  47 #include "gromacs/mdtypes/inputrec.h"
  48 #include "gromacs/mdtypes/md_enums.h"
  49 #include "gromacs/mdtypes/state.h"
  50 #include "gromacs/random/threefry.h"
  51 #include "gromacs/random/uniformrealdistribution.h"
  52 #include "gromacs/utility/enumerationhelpers.h"
  53 #include "gromacs/utility/fatalerror.h"
  54 #include "gromacs/utility/smalloc.h"
  55
  56 #include "expanded_internal.h"
  57
  58 static void init_df_history_weights(df_history_t* dfhist, const t_expanded* expand, int nlim)
  59 {
  60     int i;
  61     dfhist->wl_delta = expand->init_wl_delta;
  62     for (i = 0; i < nlim; i++)
  63     {
  64         dfhist->sum_weights[i] = expand->init_lambda_weights[i];
  65         dfhist->sum_dg[i]      = expand->init_lambda_weights[i];
  66     }
  67 }
  68
  69 /* Eventually should contain all the functions needed to initialize expanded ensemble
  70    before the md loop starts */
  71 void init_expanded_ensemble(gmx_bool bStateFromCP, const t_inputrec* ir, df_history_t* dfhist)
  72 {
  73     if (!bStateFromCP)
  74     {
  75         init_df_history_weights(dfhist, ir->expandedvals.get(), ir->fepvals->n_lambda);
  76     }
  77 }
  78
  79 static void GenerateGibbsProbabilities(const real* ene, double* p_k, double* pks, int minfep, int maxfep)
  80 {
  81
  82     int  i;
  83     real maxene;
  84
  85     *pks   = 0.0;
  86     maxene = ene[minfep];
  87     /* find the maximum value */
  88     for (i = minfep; i <= maxfep; i++)
  89     {
  90         if (ene[i] > maxene)
  91         {
  92             maxene = ene[i];
  93         }
  94     }
  95     /* find the denominator */
  96     for (i = minfep; i <= maxfep; i++)
  97     {
  98         *pks += std::exp(ene[i] - maxene);
  99     }
 100     /*numerators*/
 101     for (i = minfep; i <= maxfep; i++)
 102     {
 103         p_k[i] = std::exp(ene[i] - maxene) / *pks;
 104     }
 105 }
 106
 107 static void
 108 GenerateWeightedGibbsProbabilities(const real* ene, double* p_k, double* pks, int nlim, real* nvals, real delta)
 109 {
 110
 111     int   i;
 112     real  maxene;
 113     real* nene;
 114     *pks = 0.0;
 115
 116     snew(nene, nlim);
 117     for (i = 0; i < nlim; i++)
 118     {
 119         if (nvals[i] == 0)
 120         {
 121             /* add the delta, since we need to make sure it's greater than zero, and
 122                we need a non-arbitrary number? */
 123             nene[i] = ene[i] + std::log(nvals[i] + delta);
 124         }
 125         else
 126         {
 127             nene[i] = ene[i] + std::log(nvals[i]);
 128         }
 129     }
 130
 131     /* find the maximum value */
 132     maxene = nene[0];
 133     for (i = 0; i < nlim; i++)
 134     {
 135         if (nene[i] > maxene)
 136         {
 137             maxene = nene[i];
 138         }
 139     }
 140
 141     /* subtract off the maximum, avoiding overflow */
 142     for (i = 0; i < nlim; i++)
 143     {
 144         nene[i] -= maxene;
 145     }
 146
 147     /* find the denominator */
 148     for (i = 0; i < nlim; i++)
 149     {
 150         *pks += std::exp(nene[i]);
 151     }
 152
 153     /*numerators*/
 154     for (i = 0; i < nlim; i++)
 155     {
 156         p_k[i] = std::exp(nene[i]) / *pks;
 157     }
 158     sfree(nene);
 159 }
 160
 161 static int FindMinimum(const real* min_metric, int N)
 162 {
 163
 164     real min_val;
 165     int  min_nval, nval;
 166
 167     min_nval = 0;
 168     min_val  = min_metric[0];
 169
 170     for (nval = 0; nval < N; nval++)
 171     {
 172         if (min_metric[nval] < min_val)
 173         {
 174             min_val  = min_metric[nval];
 175             min_nval = nval;
 176         }
 177     }
 178     return min_nval;
 179 }
 180
 181 static gmx_bool CheckHistogramRatios(int nhisto, const real* histo, real ratio)
 182 {
 183
 184     int      i;
 185     real     nmean;
 186     gmx_bool bIfFlat;
 187
 188     nmean = 0;
 189     for (i = 0; i < nhisto; i++)
 190     {
 191         nmean += histo[i];
 192     }
 193
 194     if (nmean == 0)
 195     {
 196         /* no samples! is bad!*/
 197         bIfFlat = FALSE;
 198         return bIfFlat;
 199     }
 200     nmean /= static_cast<real>(nhisto);
 201
 202     bIfFlat = TRUE;
 203     for (i = 0; i < nhisto; i++)
 204     {
 205         /* make sure that all points are in the ratio < x <  1/ratio range  */
 206         if (!((histo[i] / nmean < 1.0 / ratio) && (histo[i] / nmean > ratio)))
 207         {
 208             bIfFlat = FALSE;
 209             break;
 210         }
 211     }
 212     return bIfFlat;
 213 }
 214
 215 static gmx_bool CheckIfDoneEquilibrating(int nlim, const t_expanded* expand, const df_history_t* dfhist, int64_t step)
 216 {
 217
 218     int      i, totalsamples;
 219     gmx_bool bDoneEquilibrating = TRUE;
 220     gmx_bool bIfFlat;
 221
 222     /* If we are doing slow growth to get initial values, we haven't finished equilibrating */
 223     if (expand->lmc_forced_nstart > 0)
 224     {
 225         for (i = 0; i < nlim; i++)
 226         {
 227             if (dfhist->n_at_lam[i]
 228                 < expand->lmc_forced_nstart) /* we are still doing the initial sweep, so we're
 229                                                 definitely not done equilibrating*/
 230             {
 231                 bDoneEquilibrating = FALSE;
 232                 break;
 233             }
 234         }
 235     }
 236     else
 237     {
 238         /* assume we have equilibrated the weights, then check to see if any of the conditions are not met */
 239         bDoneEquilibrating = TRUE;
 240
 241         /* calculate the total number of samples */
 242         switch (expand->elmceq)
 243         {
 244             case LambdaWeightWillReachEquilibrium::No:
 245                 /* We have not equilibrated, and won't, ever. */
 246                 bDoneEquilibrating = FALSE;
 247                 break;
 248             case LambdaWeightWillReachEquilibrium::Yes:
 249                 /* we have equilibrated -- we're done */
 250                 bDoneEquilibrating = TRUE;
 251                 break;
 252             case LambdaWeightWillReachEquilibrium::Steps:
 253                 /* first, check if we are equilibrating by steps, if we're still under */
 254                 if (step < expand->equil_steps)
 255                 {
 256                     bDoneEquilibrating = FALSE;
 257                 }
 258                 break;
 259             case LambdaWeightWillReachEquilibrium::Samples:
 260                 totalsamples = 0;
 261                 for (i = 0; i < nlim; i++)
 262                 {
 263                     totalsamples += dfhist->n_at_lam[i];
 264                 }
 265                 if (totalsamples < expand->equil_samples)
 266                 {
 267                     bDoneEquilibrating = FALSE;
 268                 }
 269                 break;
 270             case LambdaWeightWillReachEquilibrium::NumAtLambda:
 271                 for (i = 0; i < nlim; i++)
 272                 {
 273                     if (dfhist->n_at_lam[i]
 274                         < expand->equil_n_at_lam) /* we are still doing the initial sweep, so we're
 275                                                      definitely not done equilibrating*/
 276                     {
 277                         bDoneEquilibrating = FALSE;
 278                         break;
 279                     }
 280                 }
 281                 break;
 282             case LambdaWeightWillReachEquilibrium::WLDelta:
 283                 if (EWL(expand->elamstats)) /* This check is in readir as well, but
 284                                                just to be sure */
 285                 {
 286                     if (dfhist->wl_delta > expand->equil_wl_delta)
 287                     {
 288                         bDoneEquilibrating = FALSE;
 289                     }
 290                 }
 291                 break;
 292             case LambdaWeightWillReachEquilibrium::Ratio:
 293                 /* we can use the flatness as a judge of good weights, as long as
 294                    we're not doing minvar, or Wang-Landau.
 295                    But turn off for now until we figure out exactly how we do this.
 296                  */
 297
 298                 if (!(EWL(expand->elamstats) || expand->elamstats == LambdaWeightCalculation::Minvar))
 299                 {
 300                     /* we want to use flatness -avoiding- the forced-through samples.  Plus, we need
 301                        to convert to floats for this histogram function. */
 302
 303                     real* modhisto;
 304                     snew(modhisto, nlim);
 305                     for (i = 0; i < nlim; i++)
 306                     {
 307                         modhisto[i] = 1.0 * (dfhist->n_at_lam[i] - expand->lmc_forced_nstart);
 308                     }
 309                     bIfFlat = CheckHistogramRatios(nlim, modhisto, expand->equil_ratio);
 310                     sfree(modhisto);
 311                     if (!bIfFlat)
 312                     {
 313                         bDoneEquilibrating = FALSE;
 314                     }
 315                 }
 316                 break;
 317             default: bDoneEquilibrating = TRUE; break;
 318         }
 319     }
 320     return bDoneEquilibrating;
 321 }
 322
 323 static gmx_bool UpdateWeights(int           nlim,
 324                               t_expanded*   expand,
 325                               df_history_t* dfhist,
 326                               int           fep_state,
 327                               const real*   scaled_lamee,
 328                               const real*   weighted_lamee,
 329                               int64_t       step)
 330 {
 331     gmx_bool bSufficientSamples;
 332     real     acceptanceWeight;
 333     int      i;
 334     int      min_nvalm, min_nvalp, maxc;
 335     real     omega_m1_0, omega_p1_0;
 336     real     zero_sum_weights;
 337     real *omegam_array, *weightsm_array, *omegap_array, *weightsp_array, *varm_array, *varp_array,
 338             *dwp_array, *dwm_array;
 339     real    clam_varm, clam_varp, clam_osum, clam_weightsm, clam_weightsp, clam_minvar;
 340     real *  lam_variance, *lam_dg;
 341     double* p_k;
 342     double  pks = 0;
 343
 344     /* Future potential todos for this function (see #3848):
 345      *  - Update the names in the dhist structure to be clearer. Not done for now since this
 346      *    a bugfix update and we are mininizing other code changes.
 347      *  - Modularize the code some more.
 348      *  - potentially merge with accelerated weight histogram functionality, since it's very similar.
 349      */
 350     /*  if we have equilibrated the expanded ensemble weights, we are not updating them, so exit now */
 351     if (dfhist->bEquil)
 352     {
 353         return FALSE;
 354     }
 355
 356     if (CheckIfDoneEquilibrating(nlim, expand, dfhist, step))
 357     {
 358         dfhist->bEquil = TRUE;
 359         /* zero out the visited states so we know how many equilibrated states we have
 360            from here on out.*/
 361         for (i = 0; i < nlim; i++)
 362         {
 363             dfhist->n_at_lam[i] = 0;
 364         }
 365         return TRUE;
 366     }
 367
 368     /* If we reached this far, we have not equilibrated yet, keep on
 369        going resetting the weights */
 370
 371     if (EWL(expand->elamstats))
 372     {
 373         if (expand->elamstats
 374             == LambdaWeightCalculation::WL) /* Using standard Wang-Landau for weight updates */
 375         {
 376             dfhist->sum_weights[fep_state] -= dfhist->wl_delta;
 377             dfhist->wl_histo[fep_state] += 1.0;
 378         }
 379         else if (expand->elamstats == LambdaWeightCalculation::WWL)
 380         /* Using weighted Wang-Landau for weight updates.
 381          * Very closly equivalent to accelerated weight histogram approach
 382          * applied to expanded ensemble. */
 383         {
 384             snew(p_k, nlim);
 385
 386             /* first increment count */
 387             GenerateGibbsProbabilities(weighted_lamee, p_k, &pks, 0, nlim - 1);
 388             for (i = 0; i < nlim; i++)
 389             {
 390                 dfhist->wl_histo[i] += static_cast<real>(p_k[i]);
 391             }
 392
 393             /* then increment weights (uses count) */
 394             pks = 0.0;
 395             GenerateWeightedGibbsProbabilities(
 396                     weighted_lamee, p_k, &pks, nlim, dfhist->wl_histo, dfhist->wl_delta);
 397
 398             for (i = 0; i < nlim; i++)
 399             {
 400                 dfhist->sum_weights[i] -= dfhist->wl_delta * static_cast<real>(p_k[i]);
 401             }
 402             /* Alternate definition, using logarithms. Shouldn't make very much difference! */
 403             /*
 404                real di;
 405                for (i=0;i<nlim;i++)
 406                {
 407                 di = (real)1.0 + dfhist->wl_delta*(real)p_k[i];
 408                 dfhist->sum_weights[i] -= log(di);
 409                }
 410              */
 411             sfree(p_k);
 412         }
 413
 414         zero_sum_weights = dfhist->sum_weights[0];
 415         for (i = 0; i < nlim; i++)
 416         {
 417             dfhist->sum_weights[i] -= zero_sum_weights;
 418         }
 419     }
 420
 421     if (expand->elamstats == LambdaWeightCalculation::Barker
 422         || expand->elamstats == LambdaWeightCalculation::Metropolis
 423         || expand->elamstats == LambdaWeightCalculation::Minvar)
 424     {
 425         maxc = 2 * expand->c_range + 1;
 426
 427         snew(lam_dg, nlim);
 428         snew(lam_variance, nlim);
 429
 430         snew(omegap_array, maxc);
 431         snew(weightsp_array, maxc);
 432         snew(varp_array, maxc);
 433         snew(dwp_array, maxc);
 434
 435         snew(omegam_array, maxc);
 436         snew(weightsm_array, maxc);
 437         snew(varm_array, maxc);
 438         snew(dwm_array, maxc);
 439
 440         /* unpack the values of the free energy differences and the
 441          * variance in their estimates between nearby lambdas. We will
 442          * only actually update 2 of these, the state we are currently
 443          * at and the one we end up moving to
 444          */
 445
 446         for (i = 0; i < nlim - 1; i++)
 447         { /* only through the second to last */
 448             lam_dg[i] = dfhist->sum_dg[i + 1] - dfhist->sum_dg[i];
 449             lam_variance[i] =
 450                     gmx::square(dfhist->sum_variance[i + 1]) - gmx::square(dfhist->sum_variance[i]);
 451         }
 452
 453         /* accumulate running averages of thermodynamic averages for Bennett Acceptance Ratio-based
 454          * estimates of the free energy .
 455          * Rather than peforming self-consistent estimation of the free energies at each step,
 456          * we keep track of an array of possible different free energies (cnvals),
 457          * and we self-consistently choose the best one. The one that leads to a free energy estimate
 458          * that is closest to itself is the best estimate of the free energy.  It is essentially a
 459          * parallellized version of self-consistent iteration.  maxc is the number of these constants. */
 460
 461         for (int nval = 0; nval < maxc; nval++)
 462         {
 463             const real cnval = static_cast<real>(nval - expand->c_range);
 464
 465             /* Compute acceptance criterion weight to the state below this one for use in averages.
 466              * Note we do not have to have just moved from that state to use this free energy
 467              * estimate; these are essentially "virtual" moves. */
 468
 469             if (fep_state > 0)
 470             {
 471                 const auto lambdaEnergyDifference =
 472                         cnval - (scaled_lamee[fep_state] - scaled_lamee[fep_state - 1]);
 473                 acceptanceWeight =
 474                         gmx::calculateAcceptanceWeight(expand->elamstats, lambdaEnergyDifference);
 475                 dfhist->accum_m[fep_state][nval] += acceptanceWeight;
 476                 dfhist->accum_m2[fep_state][nval] += acceptanceWeight * acceptanceWeight;
 477             }
 478
 479             // Compute acceptance criterion weight to transition to the next state
 480             if (fep_state < nlim - 1)
 481             {
 482                 const auto lambdaEnergyDifference =
 483                         -cnval + (scaled_lamee[fep_state + 1] - scaled_lamee[fep_state]);
 484                 acceptanceWeight =
 485                         gmx::calculateAcceptanceWeight(expand->elamstats, lambdaEnergyDifference);
 486                 dfhist->accum_p[fep_state][nval] += acceptanceWeight;
 487                 dfhist->accum_p2[fep_state][nval] += acceptanceWeight * acceptanceWeight;
 488             }
 489
 490             /* Determination of Metropolis transition and Barker transition weights */
 491
 492             int numObservationsCurrentState = dfhist->n_at_lam[fep_state];
 493             /* determine the number of observations above and below the current state */
 494             int numObservationsLowerState = 0;
 495             if (fep_state > 0)
 496             {
 497                 numObservationsLowerState = dfhist->n_at_lam[fep_state - 1];
 498             }
 499             int numObservationsHigherState = 0;
 500             if (fep_state < nlim - 1)
 501             {
 502                 numObservationsHigherState = dfhist->n_at_lam[fep_state + 1];
 503             }
 504
 505             /* Calculate the biases for each expanded ensemble state that minimize the total
 506              * variance, as implemented in Martinez-Veracoechea and Escobedo,
 507              * J. Phys. Chem. B 2008, 112, 8120-8128
 508              *
 509              * The variance associated with the free energy estimate between two states i and j
 510              * is calculated as
 511              *     Var(i,j) = {avg[xi(i->j)^2] / avg[xi(i->j)]^2 - 1} / numObservations(i->j)
 512              *              + {avg[xi(j->i)^2] / avg[xi(j->i)]^2 - 1} / numObservations(j->i)
 513              * where xi(i->j) is the acceptance factor / weight associated with moving from state i to j
 514              * As we are calculating the acceptance factor to the neighbors every time we're visiting
 515              * a state, numObservations(i->j) == numObservations(i) and numObservations(j->i) == numObservations(j)
 516              */
 517
 518             /* Accumulation of acceptance weight averages between the current state and the
 519              * states +1 (p1) and -1 (m1), averaged at current state (0)
 520              */
 521             real avgAcceptanceCurrentToLower  = 0;
 522             real avgAcceptanceCurrentToHigher = 0;
 523             /* Accumulation of acceptance weight averages quantities between states 0
 524              *  and states +1 and -1, squared
 525              */
 526             real avgAcceptanceCurrentToLowerSquared  = 0;
 527             real avgAcceptanceCurrentToHigherSquared = 0;
 528             /* Accumulation of free energy quantities from lower state (m1) to current state (0) and squared */
 529             real avgAcceptanceLowerToCurrent        = 0;
 530             real avgAcceptanceLowerToCurrentSquared = 0;
 531             /* Accumulation of free energy quantities from upper state (p1) to current state (0) and squared */
 532             real avgAcceptanceHigherToCurrent        = 0;
 533             real avgAcceptanceHigherToCurrentSquared = 0;
 534
 535             if (numObservationsCurrentState > 0)
 536             {
 537                 avgAcceptanceCurrentToLower = dfhist->accum_m[fep_state][nval] / numObservationsCurrentState;
 538                 avgAcceptanceCurrentToHigher =
 539                         dfhist->accum_p[fep_state][nval] / numObservationsCurrentState;
 540                 avgAcceptanceCurrentToLowerSquared =
 541                         dfhist->accum_m2[fep_state][nval] / numObservationsCurrentState;
 542                 avgAcceptanceCurrentToHigherSquared =
 543                         dfhist->accum_p2[fep_state][nval] / numObservationsCurrentState;
 544             }
 545
 546             if ((fep_state > 0) && (numObservationsLowerState > 0))
 547             {
 548                 avgAcceptanceLowerToCurrent =
 549                         dfhist->accum_p[fep_state - 1][nval] / numObservationsLowerState;
 550                 avgAcceptanceLowerToCurrentSquared =
 551                         dfhist->accum_p2[fep_state - 1][nval] / numObservationsLowerState;
 552             }
 553
 554             if ((fep_state < nlim - 1) && (numObservationsHigherState > 0))
 555             {
 556                 avgAcceptanceHigherToCurrent =
 557                         dfhist->accum_m[fep_state + 1][nval] / numObservationsHigherState;
 558                 avgAcceptanceHigherToCurrentSquared =
 559                         dfhist->accum_m2[fep_state + 1][nval] / numObservationsHigherState;
 560             }
 561             /* These are accumulation of positive values (see definition of acceptance functions
 562              * above), or of squares of positive values.
 563              * We're taking this for granted in the following calculation, so make sure
 564              * here that nothing weird happened. Although technically all values should be positive,
 565              * because of floating point precisions, they might be numerically zero. */
 566             GMX_RELEASE_ASSERT(
 567                     avgAcceptanceCurrentToLower >= 0 && avgAcceptanceCurrentToLowerSquared >= 0
 568                             && avgAcceptanceCurrentToHigher >= 0
 569                             && avgAcceptanceCurrentToHigherSquared >= 0 && avgAcceptanceLowerToCurrent >= 0
 570                             && avgAcceptanceLowerToCurrentSquared >= 0 && avgAcceptanceHigherToCurrent >= 0
 571                             && avgAcceptanceHigherToCurrentSquared >= 0,
 572                     "By definition, the acceptance factors should all be nonnegative.");
 573
 574             real varianceCurrentToLower   = 0;
 575             real varianceCurrentToHigher  = 0;
 576             real weightDifferenceToLower  = 0;
 577             real weightDifferenceToHigher = 0;
 578             real varianceToLower          = 0;
 579             real varianceToHigher         = 0;
 580
 581             if (fep_state > 0)
 582             {
 583                 if (numObservationsCurrentState > 0)
 584                 {
 585                     /* Calculate {avg[xi(i->j)^2] / avg[xi(i->j)]^2 - 1}
 586                      *
 587                      * Note that if avg[xi(i->j)] == 0, also avg[xi(i->j)^2] == 0 (since the
 588                      * acceptances are all positive!), and hence
 589                      *     {avg[xi(i->j)^2] / avg[xi(i->j)]^2 - 1} -> 0  for  avg[xi(i->j)] -> 0
 590                      * We're catching that case explicitly to avoid numerical
 591                      * problems dividing by zero when the overlap between states is small (#3304)
 592                      */
 593                     if (avgAcceptanceCurrentToLower > 0)
 594                     {
 595                         varianceCurrentToLower =
 596                                 avgAcceptanceCurrentToLowerSquared
 597                                         / (avgAcceptanceCurrentToLower * avgAcceptanceCurrentToLower)
 598                                 - 1.0;
 599                     }
 600                     if (numObservationsLowerState > 0)
 601                     {
 602                         /* Calculate {avg[xi(i->j)^2] / avg[xi(i->j)]^2 - 1}
 603                          *
 604                          * Note that if avg[xi(i->j)] == 0, also avg[xi(i->j)^2] == 0 (since the
 605                          * acceptances are all positive!), and hence
 606                          *     {avg[xi(i->j)^2] / avg[xi(i->j)]^2 - 1} -> 0  for  avg[xi(i->j)] -> 0
 607                          * We're catching that case explicitly to avoid numerical
 608                          * problems dividing by zero when the overlap between states is small (#3304)
 609                          */
 610                         real varianceLowerToCurrent = 0;
 611                         if (avgAcceptanceLowerToCurrent > 0)
 612                         {
 613                             varianceLowerToCurrent =
 614                                     avgAcceptanceLowerToCurrentSquared
 615                                             / (avgAcceptanceLowerToCurrent * avgAcceptanceLowerToCurrent)
 616                                     - 1.0;
 617                         }
 618                         /* Free energy difference to the state one state lower */
 619                         /* if these either of these quantities are zero, the energies are */
 620                         /* way too large for the dynamic range.  We need an alternate guesstimate */
 621                         if ((avgAcceptanceCurrentToLower == 0) || (avgAcceptanceLowerToCurrent == 0))
 622                         {
 623                             weightDifferenceToLower =
 624                                     (scaled_lamee[fep_state] - scaled_lamee[fep_state - 1]);
 625                         }
 626                         else
 627                         {
 628                             weightDifferenceToLower = (std::log(avgAcceptanceCurrentToLower)
 629                                                        - std::log(avgAcceptanceLowerToCurrent))
 630                                                       + cnval;
 631                         }
 632                         /* Variance of the free energy difference to the one state lower */
 633                         varianceToLower =
 634                                 (1.0 / numObservationsCurrentState) * (varianceCurrentToLower)
 635                                 + (1.0 / numObservationsLowerState) * (varianceLowerToCurrent);
 636                     }
 637                 }
 638             }
 639
 640             if (fep_state < nlim - 1)
 641             {
 642                 if (numObservationsCurrentState > 0)
 643                 {
 644                     /* Calculate {avg[xi(i->j)^2] / avg[xi(i->j)]^2 - 1}
 645                      *
 646                      * Note that if avg[xi(i->j)] == 0, also avg[xi(i->j)^2] == 0 (since the
 647                      * acceptances are all positive!), and hence
 648                      *     {avg[xi(i->j)^2] / avg[xi(i->j)]^2 - 1} -> 0  for  avg[xi(i->j)] -> 0
 649                      * We're catching that case explicitly to avoid numerical
 650                      * problems dividing by zero when the overlap between states is small (#3304)
 651                      */
 652
 653                     if (avgAcceptanceCurrentToHigher < 0)
 654                     {
 655                         varianceCurrentToHigher =
 656                                 avgAcceptanceCurrentToHigherSquared
 657                                         / (avgAcceptanceCurrentToHigher * avgAcceptanceCurrentToHigher)
 658                                 - 1.0;
 659                     }
 660                     if (numObservationsHigherState > 0)
 661                     {
 662                         /* Calculate {avg[xi(i->j)^2] / avg[xi(i->j)]^2 - 1}
 663                          *
 664                          * Note that if avg[xi(i->j)] == 0, also avg[xi(i->j)^2] == 0 (since the
 665                          * acceptances are all positive!), and hence
 666                          *     {avg[xi(i->j)^2] / avg[xi(i->j)]^2 - 1} -> 0  for  avg[xi(i->j)] -> 0
 667                          * We're catching that case explicitly to avoid numerical
 668                          * problems dividing by zero when the overlap between states is small (#3304)
 669                          */
 670                         real varianceHigherToCurrent = 0;
 671                         if (avgAcceptanceHigherToCurrent > 0)
 672                         {
 673                             varianceHigherToCurrent =
 674                                     avgAcceptanceHigherToCurrentSquared
 675                                             / (avgAcceptanceHigherToCurrent * avgAcceptanceHigherToCurrent)
 676                                     - 1.0;
 677                         }
 678                         /* Free energy difference to the state one state higher */
 679                         /* if these either of these quantities are zero, the energies are */
 680                         /* way too large for the dynamic range.  We need an alternate guesstimate */
 681                         if ((avgAcceptanceHigherToCurrent == 0) || (avgAcceptanceCurrentToHigher == 0))
 682                         {
 683                             weightDifferenceToHigher =
 684                                     (scaled_lamee[fep_state + 1] - scaled_lamee[fep_state]);
 685                         }
 686                         else
 687                         {
 688                             weightDifferenceToHigher = (std::log(avgAcceptanceHigherToCurrent)
 689                                                         - std::log(avgAcceptanceCurrentToHigher))
 690                                                        + cnval;
 691                         }
 692                         /* Variance of the free energy difference to the one state higher */
 693                         varianceToHigher =
 694                                 (1.0 / numObservationsHigherState) * (varianceHigherToCurrent)
 695                                 + (1.0 / numObservationsCurrentState) * (varianceCurrentToHigher);
 696                     }
 697                 }
 698             }
 699
 700             if (numObservationsCurrentState > 0)
 701             {
 702                 omegam_array[nval] = varianceCurrentToLower;
 703             }
 704             else
 705             {
 706                 omegam_array[nval] = 0;
 707             }
 708             weightsm_array[nval] = weightDifferenceToLower;
 709             varm_array[nval]     = varianceToLower;
 710             if (numObservationsLowerState > 0)
 711             {
 712                 dwm_array[nval] =
 713                         fabs((cnval + std::log((1.0 * numObservationsCurrentState) / numObservationsLowerState))
 714                              - lam_dg[fep_state - 1]);
 715             }
 716             else
 717             {
 718                 dwm_array[nval] = std::fabs(cnval - lam_dg[fep_state - 1]);
 719             }
 720
 721             if (numObservationsCurrentState > 0)
 722             {
 723                 omegap_array[nval] = varianceCurrentToHigher;
 724             }
 725             else
 726             {
 727                 omegap_array[nval] = 0;
 728             }
 729             weightsp_array[nval] = weightDifferenceToHigher;
 730             varp_array[nval]     = varianceToHigher;
 731             if ((numObservationsHigherState > 0) && (numObservationsCurrentState > 0))
 732             {
 733                 dwp_array[nval] =
 734                         fabs((cnval + std::log((1.0 * numObservationsHigherState) / numObservationsCurrentState))
 735                              - lam_dg[fep_state]);
 736             }
 737             else
 738             {
 739                 dwp_array[nval] = std::fabs(cnval - lam_dg[fep_state]);
 740             }
 741         }
 742
 743         /* find the free energy estimate closest to the guessed weight's value */
 744
 745         min_nvalm     = FindMinimum(dwm_array, maxc);
 746         omega_m1_0    = omegam_array[min_nvalm];
 747         clam_weightsm = weightsm_array[min_nvalm];
 748         clam_varm     = varm_array[min_nvalm];
 749
 750         min_nvalp     = FindMinimum(dwp_array, maxc);
 751         omega_p1_0    = omegap_array[min_nvalp];
 752         clam_weightsp = weightsp_array[min_nvalp];
 753         clam_varp     = varp_array[min_nvalp];
 754
 755         clam_osum   = omega_m1_0 + omega_p1_0;
 756         clam_minvar = 0;
 757         if (clam_osum > 0)
 758         {
 759             clam_minvar = 0.5 * std::log(clam_osum);
 760         }
 761
 762         if (fep_state > 0)
 763         {
 764             lam_dg[fep_state - 1]       = clam_weightsm;
 765             lam_variance[fep_state - 1] = clam_varm;
 766         }
 767
 768         if (fep_state < nlim - 1)
 769         {
 770             lam_dg[fep_state]       = clam_weightsp;
 771             lam_variance[fep_state] = clam_varp;
 772         }
 773
 774         if (expand->elamstats == LambdaWeightCalculation::Minvar)
 775         {
 776             bSufficientSamples = TRUE;
 777             /* make sure the number of samples in each state are all
 778              * past a user-specified threshold
 779              */
 780             for (i = 0; i < nlim; i++)
 781             {
 782                 if (dfhist->n_at_lam[i] < expand->minvarmin)
 783                 {
 784                     bSufficientSamples = FALSE;
 785                 }
 786             }
 787             if (bSufficientSamples)
 788             {
 789                 dfhist->sum_minvar[fep_state] = clam_minvar;
 790                 if (fep_state == 0)
 791                 {
 792                     for (i = 0; i < nlim; i++)
 793                     {
 794                         dfhist->sum_minvar[i] += (expand->minvar_const - clam_minvar);
 795                     }
 796                     expand->minvar_const          = clam_minvar;
 797                     dfhist->sum_minvar[fep_state] = 0.0;
 798                 }
 799                 else
 800                 {
 801                     dfhist->sum_minvar[fep_state] -= expand->minvar_const;
 802                 }
 803             }
 804         }
 805
 806         /* we need to rezero minvar now, since it could change at fep_state = 0 */
 807         dfhist->sum_dg[0]       = 0.0;
 808         dfhist->sum_variance[0] = 0.0;
 809         dfhist->sum_weights[0]  = dfhist->sum_dg[0] + dfhist->sum_minvar[0]; /* should be zero */
 810
 811         for (i = 1; i < nlim; i++)
 812         {
 813             dfhist->sum_dg[i] = lam_dg[i - 1] + dfhist->sum_dg[i - 1];
 814             dfhist->sum_variance[i] =
 815                     std::sqrt(lam_variance[i - 1] + gmx::square(dfhist->sum_variance[i - 1]));
 816             dfhist->sum_weights[i] = dfhist->sum_dg[i] + dfhist->sum_minvar[i];
 817         }
 818
 819         sfree(lam_dg);
 820         sfree(lam_variance);
 821
 822         sfree(omegam_array);
 823         sfree(weightsm_array);
 824         sfree(varm_array);
 825         sfree(dwm_array);
 826
 827         sfree(omegap_array);
 828         sfree(weightsp_array);
 829         sfree(varp_array);
 830         sfree(dwp_array);
 831     }
 832     return FALSE;
 833 }
 834
 835 static int ChooseNewLambda(int               nlim,
 836                            const t_expanded* expand,
 837                            df_history_t*     dfhist,
 838                            int               fep_state,
 839                            const real*       weighted_lamee,
 840                            double*           p_k,
 841                            int64_t           seed,
 842                            int64_t           step)
 843 {
 844     /* Choose new lambda value, and update transition matrix */
 845
 846     int                  i, ifep, minfep, maxfep, lamnew, lamtrial, starting_fep_state;
 847     real                 r1, r2, de, trialprob, tprob = 0;
 848     double *             propose, *accept, *remainder;
 849     double               pks;
 850     real                 pnorm;
 851     gmx::ThreeFry2x64<0> rng(
 852             seed, gmx::RandomDomain::ExpandedEnsemble); // We only draw once, so zero bits internal counter is fine
 853     gmx::UniformRealDistribution<real> dist;
 854
 855     starting_fep_state = fep_state;
 856     lamnew             = fep_state; /* so that there is a default setting -- stays the same */
 857
 858     if (!EWL(expand->elamstats)) /* ignore equilibrating the weights if using WL */
 859     {
 860         if ((expand->lmc_forced_nstart > 0) && (dfhist->n_at_lam[nlim - 1] <= expand->lmc_forced_nstart))
 861         {
 862             /* Use a marching method to run through the lambdas and get preliminary free energy data,
 863                before starting 'free' sampling.  We start free sampling when we have enough at each lambda */
 864
 865             /* if we have enough at this lambda, move on to the next one */
 866
 867             if (dfhist->n_at_lam[fep_state] == expand->lmc_forced_nstart)
 868             {
 869                 lamnew = fep_state + 1;
 870                 if (lamnew == nlim) /* whoops, stepped too far! */
 871                 {
 872                     lamnew -= 1;
 873                 }
 874             }
 875             else
 876             {
 877                 lamnew = fep_state;
 878             }
 879             return lamnew;
 880         }
 881     }
 882
 883     snew(propose, nlim);
 884     snew(accept, nlim);
 885     snew(remainder, nlim);
 886
 887     for (i = 0; i < expand->lmc_repeats; i++)
 888     {
 889         rng.restart(step, i);
 890         dist.reset();
 891
 892         for (ifep = 0; ifep < nlim; ifep++)
 893         {
 894             propose[ifep] = 0;
 895             accept[ifep]  = 0;
 896         }
 897
 898         if ((expand->elmcmove == LambdaMoveCalculation::Gibbs)
 899             || (expand->elmcmove == LambdaMoveCalculation::MetropolisGibbs))
 900         {
 901             /* use the Gibbs sampler, with restricted range */
 902             if (expand->gibbsdeltalam < 0)
 903             {
 904                 minfep = 0;
 905                 maxfep = nlim - 1;
 906             }
 907             else
 908             {
 909                 minfep = fep_state - expand->gibbsdeltalam;
 910                 maxfep = fep_state + expand->gibbsdeltalam;
 911                 if (minfep < 0)
 912                 {
 913                     minfep = 0;
 914                 }
 915                 if (maxfep > nlim - 1)
 916                 {
 917                     maxfep = nlim - 1;
 918                 }
 919             }
 920
 921             GenerateGibbsProbabilities(weighted_lamee, p_k, &pks, minfep, maxfep);
 922
 923             if (expand->elmcmove == LambdaMoveCalculation::Gibbs)
 924             {
 925                 for (ifep = minfep; ifep <= maxfep; ifep++)
 926                 {
 927                     propose[ifep] = p_k[ifep];
 928                     accept[ifep]  = 1.0;
 929                 }
 930                 /* Gibbs sampling */
 931                 r1 = dist(rng);
 932                 for (lamnew = minfep; lamnew <= maxfep; lamnew++)
 933                 {
 934                     if (r1 <= p_k[lamnew])
 935                     {
 936                         break;
 937                     }
 938                     r1 -= p_k[lamnew];
 939                 }
 940             }
 941             else if (expand->elmcmove == LambdaMoveCalculation::MetropolisGibbs)
 942             {
 943
 944                 /* Metropolized Gibbs sampling */
 945                 for (ifep = minfep; ifep <= maxfep; ifep++)
 946                 {
 947                     remainder[ifep] = 1 - p_k[ifep];
 948                 }
 949
 950                 /* find the proposal probabilities */
 951
 952                 if (remainder[fep_state] == 0)
 953                 {
 954                     /* only the current state has any probability */
 955                     /* we have to stay at the current state */
 956                     lamnew = fep_state;
 957                 }
 958                 else
 959                 {
 960                     for (ifep = minfep; ifep <= maxfep; ifep++)
 961                     {
 962                         if (ifep != fep_state)
 963                         {
 964                             propose[ifep] = p_k[ifep] / remainder[fep_state];
 965                         }
 966                         else
 967                         {
 968                             propose[ifep] = 0;
 969                         }
 970                     }
 971
 972                     r1 = dist(rng);
 973                     for (lamtrial = minfep; lamtrial <= maxfep; lamtrial++)
 974                     {
 975                         pnorm = p_k[lamtrial] / remainder[fep_state];
 976                         if (lamtrial != fep_state)
 977                         {
 978                             if (r1 <= pnorm)
 979                             {
 980                                 break;
 981                             }
 982                             r1 -= pnorm;
 983                         }
 984                     }
 985
 986                     /* we have now selected lamtrial according to p(lamtrial)/1-p(fep_state) */
 987                     tprob = 1.0;
 988                     /* trial probability is min{1,\frac{1 - p(old)}{1-p(new)} MRS 1/8/2008 */
 989                     trialprob = (remainder[fep_state]) / (remainder[lamtrial]);
 990                     if (trialprob < tprob)
 991                     {
 992                         tprob = trialprob;
 993                     }
 994                     r2 = dist(rng);
 995                     if (r2 < tprob)
 996                     {
 997                         lamnew = lamtrial;
 998                     }
 999                     else
1000                     {
1001                         lamnew = fep_state;
1002                     }
1003                 }
1004
1005                 /* now figure out the acceptance probability for each */
1006                 for (ifep = minfep; ifep <= maxfep; ifep++)
1007                 {
1008                     tprob = 1.0;
1009                     if (remainder[ifep] != 0)
1010                     {
1011                         trialprob = (remainder[fep_state]) / (remainder[ifep]);
1012                     }
1013                     else
1014                     {
1015                         trialprob = 1.0; /* this state is the only choice! */
1016                     }
1017                     if (trialprob < tprob)
1018                     {
1019                         tprob = trialprob;
1020                     }
1021                     /* probability for fep_state=0, but that's fine, it's never proposed! */
1022                     accept[ifep] = tprob;
1023                 }
1024             }
1025
1026             if (lamnew > maxfep)
1027             {
1028                 /* it's possible some rounding is failing */
1029                 if (gmx_within_tol(remainder[fep_state], 0, 50 * GMX_DOUBLE_EPS))
1030                 {
1031                     /* numerical rounding error -- no state other than the original has weight */
1032                     lamnew = fep_state;
1033                 }
1034                 else
1035                 {
1036                     /* probably not a numerical issue */
1037                     int   loc    = 0;
1038                     int   nerror = 200 + (maxfep - minfep + 1) * 60;
1039                     char* errorstr;
1040                     snew(errorstr, nerror);
1041                     /* if its greater than maxfep, then something went wrong -- probably underflow
1042                        in the calculation of sum weights. Generated detailed info for failure */
1043                     loc += sprintf(
1044                             errorstr,
1045                             "Something wrong in choosing new lambda state with a Gibbs move -- "
1046                             "probably underflow in weight determination.\nDenominator is: "
1047                             "%3d%17.10e\n  i                dE        numerator          weights\n",
1048                             0,
1049                             pks);
1050                     for (ifep = minfep; ifep <= maxfep; ifep++)
1051                     {
1052                         loc += sprintf(&errorstr[loc],
1053                                        "%3d %17.10e%17.10e%17.10e\n",
1054                                        ifep,
1055                                        weighted_lamee[ifep],
1056                                        p_k[ifep],
1057                                        dfhist->sum_weights[ifep]);
1058                     }
1059                     gmx_fatal(FARGS, "%s", errorstr);
1060                 }
1061             }
1062         }
1063         else if ((expand->elmcmove == LambdaMoveCalculation::Metropolis)
1064                  || (expand->elmcmove == LambdaMoveCalculation::Barker))
1065         {
1066             /* use the metropolis sampler with trial +/- 1 */
1067             r1 = dist(rng);
1068             if (r1 < 0.5)
1069             {
1070                 if (fep_state == 0)
1071                 {
1072                     lamtrial = fep_state;
1073                 }
1074                 else
1075                 {
1076                     lamtrial = fep_state - 1;
1077                 }
1078             }
1079             else
1080             {
1081                 if (fep_state == nlim - 1)
1082                 {
1083                     lamtrial = fep_state;
1084                 }
1085                 else
1086                 {
1087                     lamtrial = fep_state + 1;
1088                 }
1089             }
1090
1091             de = weighted_lamee[lamtrial] - weighted_lamee[fep_state];
1092             if (expand->elmcmove == LambdaMoveCalculation::Metropolis)
1093             {
1094                 tprob = 1.0;
1095                 if (de < 0)
1096                 {
1097                     tprob = std::exp(de);
1098                 }
1099                 propose[fep_state] = 0;
1100                 propose[lamtrial]  = 1.0; /* note that this overwrites the above line if fep_state = ntrial, which only occurs at the ends */
1101                 accept[fep_state] =
1102                         1.0; /* doesn't actually matter, never proposed unless fep_state = ntrial, in which case it's 1.0 anyway */
1103                 accept[lamtrial] = tprob;
1104             }
1105             else if (expand->elmcmove == LambdaMoveCalculation::Barker)
1106             {
1107                 if (de > 0) /* Numerically stable version */
1108                 {
1109                     tprob = 1.0 / (1.0 + std::exp(-de));
1110                 }
1111                 else if (de < 0)
1112                 {
1113                     tprob = std::exp(de) / (std::exp(de) + 1.0);
1114                 }
1115                 propose[fep_state] = (1 - tprob);
1116                 propose[lamtrial] +=
1117                         tprob; /* we add, to account for the fact that at the end, they might be the same point */
1118                 accept[fep_state] = 1.0;
1119                 accept[lamtrial]  = 1.0;
1120             }
1121
1122             r2 = dist(rng);
1123             if (r2 < tprob)
1124             {
1125                 lamnew = lamtrial;
1126             }
1127             else
1128             {
1129                 lamnew = fep_state;
1130             }
1131         }
1132
1133         for (ifep = 0; ifep < nlim; ifep++)
1134         {
1135             dfhist->Tij[fep_state][ifep] += propose[ifep] * accept[ifep];
1136             dfhist->Tij[fep_state][fep_state] += propose[ifep] * (1.0 - accept[ifep]);
1137         }
1138         fep_state = lamnew;
1139     }
1140
1141     dfhist->Tij_empirical[starting_fep_state][lamnew] += 1.0;
1142
1143     sfree(propose);
1144     sfree(accept);
1145     sfree(remainder);
1146
1147     return lamnew;
1148 }
1149
1150 /* print out the weights to the log, along with current state */
1151 void PrintFreeEnergyInfoToFile(FILE*               outfile,
1152                                const t_lambda*     fep,
1153                                const t_expanded*   expand,
1154                                const t_simtemp*    simtemp,
1155                                const df_history_t* dfhist,
1156                                int                 fep_state,
1157                                int                 frequency,
1158                                int64_t             step)
1159 {
1160     int      nlim, ifep, jfep;
1161     real     dw, dg, dv, Tprint;
1162     gmx_bool bSimTemp = FALSE;
1163
1164     nlim = fep->n_lambda;
1165     if (simtemp != nullptr)
1166     {
1167         bSimTemp = TRUE;
1168     }
1169
1170     if (step % frequency == 0)
1171     {
1172         fprintf(outfile, "             MC-lambda information\n");
1173         if (EWL(expand->elamstats) && (!(dfhist->bEquil)))
1174         {
1175             fprintf(outfile, "  Wang-Landau incrementor is: %11.5g\n", dfhist->wl_delta);
1176         }
1177         fprintf(outfile, "  N");
1178         for (auto i : keysOf(fep->separate_dvdl))
1179         {
1180             if (fep->separate_dvdl[i])
1181             {
1182                 fprintf(outfile, "%7s", enumValueToString(i));
1183             }
1184             else if ((i == FreeEnergyPerturbationCouplingType::Temperature) && bSimTemp)
1185             {
1186                 fprintf(outfile, "%10s", enumValueToString(i)); /* more space for temperature formats */
1187             }
1188         }
1189         fprintf(outfile, "    Count   ");
1190         if (expand->elamstats == LambdaWeightCalculation::Minvar)
1191         {
1192             fprintf(outfile, "W(in kT)   G(in kT)  dG(in kT)  dV(in kT)\n");
1193         }
1194         else
1195         {
1196             fprintf(outfile, "G(in kT)  dG(in kT)\n");
1197         }
1198         for (ifep = 0; ifep < nlim; ifep++)
1199         {
1200             if (ifep == nlim - 1)
1201             {
1202                 dw = 0.0;
1203                 dg = 0.0;
1204                 dv = 0.0;
1205             }
1206             else
1207             {
1208                 dw = dfhist->sum_weights[ifep + 1] - dfhist->sum_weights[ifep];
1209                 dg = dfhist->sum_dg[ifep + 1] - dfhist->sum_dg[ifep];
1210                 dv = std::sqrt(gmx::square(dfhist->sum_variance[ifep + 1])
1211                                - gmx::square(dfhist->sum_variance[ifep]));
1212             }
1213             fprintf(outfile, "%3d", (ifep + 1));
1214             for (auto i : keysOf(fep->separate_dvdl))
1215             {
1216                 if (fep->separate_dvdl[i])
1217                 {
1218                     fprintf(outfile, "%7.3f", fep->all_lambda[i][ifep]);
1219                 }
1220                 else if (i == FreeEnergyPerturbationCouplingType::Temperature && bSimTemp)
1221                 {
1222                     fprintf(outfile, "%9.3f", simtemp->temperatures[ifep]);
1223                 }
1224             }
1225             if (EWL(expand->elamstats)
1226                 && (!(dfhist->bEquil))) /* if performing WL and still haven't equilibrated */
1227             {
1228                 if (expand->elamstats == LambdaWeightCalculation::WL)
1229                 {
1230                     fprintf(outfile, " %8d", static_cast<int>(dfhist->wl_histo[ifep]));
1231                 }
1232                 else
1233                 {
1234                     fprintf(outfile, " %8.3f", dfhist->wl_histo[ifep]);
1235                 }
1236             }
1237             else /* we have equilibrated weights */
1238             {
1239                 fprintf(outfile, " %8d", dfhist->n_at_lam[ifep]);
1240             }
1241             if (expand->elamstats == LambdaWeightCalculation::Minvar)
1242             {
1243                 fprintf(outfile,
1244                         " %10.5f %10.5f %10.5f %10.5f",
1245                         dfhist->sum_weights[ifep],
1246                         dfhist->sum_dg[ifep],
1247                         dg,
1248                         dv);
1249             }
1250             else
1251             {
1252                 fprintf(outfile, " %10.5f %10.5f", dfhist->sum_weights[ifep], dw);
1253             }
1254             if (ifep == fep_state)
1255             {
1256                 fprintf(outfile, " <<\n");
1257             }
1258             else
1259             {
1260                 fprintf(outfile, "   \n");
1261             }
1262         }
1263         fprintf(outfile, "\n");
1264
1265         if ((step % expand->nstTij == 0) && (expand->nstTij > 0) && (step > 0))
1266         {
1267             fprintf(outfile, "                     Transition Matrix\n");
1268             for (ifep = 0; ifep < nlim; ifep++)
1269             {
1270                 fprintf(outfile, "%12d", (ifep + 1));
1271             }
1272             fprintf(outfile, "\n");
1273             for (ifep = 0; ifep < nlim; ifep++)
1274             {
1275                 for (jfep = 0; jfep < nlim; jfep++)
1276                 {
1277                     if (dfhist->n_at_lam[ifep] > 0)
1278                     {
1279                         if (expand->bSymmetrizedTMatrix)
1280                         {
1281                             Tprint = (dfhist->Tij[ifep][jfep] + dfhist->Tij[jfep][ifep])
1282                                      / (dfhist->n_at_lam[ifep] + dfhist->n_at_lam[jfep]);
1283                         }
1284                         else
1285                         {
1286                             Tprint = (dfhist->Tij[ifep][jfep]) / (dfhist->n_at_lam[ifep]);
1287                         }
1288                     }
1289                     else
1290                     {
1291                         Tprint = 0.0;
1292                     }
1293                     fprintf(outfile, "%12.8f", Tprint);
1294                 }
1295                 fprintf(outfile, "%3d\n", (ifep + 1));
1296             }
1297
1298             fprintf(outfile, "                  Empirical Transition Matrix\n");
1299             for (ifep = 0; ifep < nlim; ifep++)
1300             {
1301                 fprintf(outfile, "%12d", (ifep + 1));
1302             }
1303             fprintf(outfile, "\n");
1304             for (ifep = 0; ifep < nlim; ifep++)
1305             {
1306                 for (jfep = 0; jfep < nlim; jfep++)
1307                 {
1308                     if (dfhist->n_at_lam[ifep] > 0)
1309                     {
1310                         if (expand->bSymmetrizedTMatrix)
1311                         {
1312                             Tprint = (dfhist->Tij_empirical[ifep][jfep] + dfhist->Tij_empirical[jfep][ifep])
1313                                      / (dfhist->n_at_lam[ifep] + dfhist->n_at_lam[jfep]);
1314                         }
1315                         else
1316                         {
1317                             Tprint = dfhist->Tij_empirical[ifep][jfep] / (dfhist->n_at_lam[ifep]);
1318                         }
1319                     }
1320                     else
1321                     {
1322                         Tprint = 0.0;
1323                     }
1324                     fprintf(outfile, "%12.8f", Tprint);
1325                 }
1326                 fprintf(outfile, "%3d\n", (ifep + 1));
1327             }
1328         }
1329     }
1330 }
1331
1332 int expandedEnsembleUpdateLambdaState(FILE*                 log,
1333                                       const t_inputrec*     ir,
1334                                       const gmx_enerdata_t* enerd,
1335                                       int                   fep_state,
1336                                       df_history_t*         dfhist,
1337                                       int64_t               step)
1338 {
1339     real *      pfep_lamee, *scaled_lamee, *weighted_lamee;
1340     double*     p_k;
1341     int         i, nlim, lamnew, totalsamples;
1342     real        oneovert, maxscaled = 0, maxweighted = 0;
1343     t_expanded* expand;
1344     t_simtemp*  simtemp;
1345     gmx_bool    bIfReset, bSwitchtoOneOverT, bDoneEquilibrating = FALSE;
1346
1347     expand  = ir->expandedvals.get();
1348     simtemp = ir->simtempvals.get();
1349     nlim    = ir->fepvals->n_lambda;
1350
1351     snew(scaled_lamee, nlim);
1352     snew(weighted_lamee, nlim);
1353     snew(pfep_lamee, nlim);
1354     snew(p_k, nlim);
1355
1356     /* update the count at the current lambda*/
1357     dfhist->n_at_lam[fep_state]++;
1358
1359     /* need to calculate the PV term somewhere, but not needed here? Not until there's a lambda
1360        state that's pressure controlled.*/
1361     /*
1362        pVTerm = 0;
1363        where does this PV term go?
1364        for (i=0;i<nlim;i++)
1365        {
1366        fep_lamee[i] += pVTerm;
1367        }
1368      */
1369
1370     /* determine the minimum value to avoid overflow.  Probably a better way to do this */
1371     /* we don't need to include the pressure term, since the volume is the same between the two.
1372        is there some term we are neglecting, however? */
1373
1374     if (ir->efep != FreeEnergyPerturbationType::No)
1375     {
1376         for (i = 0; i < nlim; i++)
1377         {
1378             if (ir->bSimTemp)
1379             {
1380                 /* Note -- this assumes no mass changes, since kinetic energy is not added  . . . */
1381                 scaled_lamee[i] =
1382                         enerd->foreignLambdaTerms.deltaH(i) / (simtemp->temperatures[i] * gmx::c_boltz)
1383                         + enerd->term[F_EPOT]
1384                                   * (1.0 / (simtemp->temperatures[i])
1385                                      - 1.0 / (simtemp->temperatures[fep_state]))
1386                                   / gmx::c_boltz;
1387             }
1388             else
1389             {
1390                 scaled_lamee[i] = enerd->foreignLambdaTerms.deltaH(i) / (expand->mc_temp * gmx::c_boltz);
1391                 /* mc_temp is currently set to the system reft unless otherwise defined */
1392             }
1393
1394             /* save these energies for printing, so they don't get overwritten by the next step */
1395             /* they aren't overwritten in the non-free energy case, but we always print with these
1396                for simplicity */
1397         }
1398     }
1399     else
1400     {
1401         if (ir->bSimTemp)
1402         {
1403             for (i = 0; i < nlim; i++)
1404             {
1405                 scaled_lamee[i] =
1406                         enerd->term[F_EPOT]
1407                         * (1.0 / simtemp->temperatures[i] - 1.0 / simtemp->temperatures[fep_state])
1408                         / gmx::c_boltz;
1409             }
1410         }
1411     }
1412
1413     for (i = 0; i < nlim; i++)
1414     {
1415         pfep_lamee[i] = scaled_lamee[i];
1416
1417         weighted_lamee[i] = dfhist->sum_weights[i] - scaled_lamee[i];
1418         if (i == 0)
1419         {
1420             maxscaled   = scaled_lamee[i];
1421             maxweighted = weighted_lamee[i];
1422         }
1423         else
1424         {
1425             if (scaled_lamee[i] > maxscaled)
1426             {
1427                 maxscaled = scaled_lamee[i];
1428             }
1429             if (weighted_lamee[i] > maxweighted)
1430             {
1431                 maxweighted = weighted_lamee[i];
1432             }
1433         }
1434     }
1435
1436     for (i = 0; i < nlim; i++)
1437     {
1438         scaled_lamee[i] -= maxscaled;
1439         weighted_lamee[i] -= maxweighted;
1440     }
1441
1442     /* update weights - we decide whether or not to actually do this inside */
1443
1444     bDoneEquilibrating =
1445             UpdateWeights(nlim, expand, dfhist, fep_state, scaled_lamee, weighted_lamee, step);
1446     if (bDoneEquilibrating)
1447     {
1448         if (log)
1449         {
1450             fprintf(log,
1451                     "\nStep %" PRId64 ": Weights have equilibrated, using criteria: %s\n",
1452                     step,
1453                     enumValueToString(expand->elmceq));
1454         }
1455     }
1456
1457     lamnew = ChooseNewLambda(
1458             nlim, expand, dfhist, fep_state, weighted_lamee, p_k, ir->expandedvals->lmc_seed, step);
1459
1460     /* now check on the Wang-Landau updating critera */
1461
1462     if (EWL(expand->elamstats))
1463     {
1464         bSwitchtoOneOverT = FALSE;
1465         if (expand->bWLoneovert)
1466         {
1467             totalsamples = 0;
1468             for (i = 0; i < nlim; i++)
1469             {
1470                 totalsamples += dfhist->n_at_lam[i];
1471             }
1472             oneovert = (1.0 * nlim) / totalsamples;
1473             /* oneovert has decreasd by a bit since last time, so we actually make sure its within one of this number */
1474             /* switch to 1/t incrementing when wl_delta has decreased at least once, and wl_delta is now less than 1/t */
1475             if ((dfhist->wl_delta <= ((totalsamples) / (totalsamples - 1.00001)) * oneovert)
1476                 && (dfhist->wl_delta < expand->init_wl_delta))
1477             {
1478                 bSwitchtoOneOverT = TRUE;
1479             }
1480         }
1481         if (bSwitchtoOneOverT)
1482         {
1483             dfhist->wl_delta =
1484                     oneovert; /* now we reduce by this each time, instead of only at flatness */
1485         }
1486         else
1487         {
1488             bIfReset = CheckHistogramRatios(nlim, dfhist->wl_histo, expand->wl_ratio);
1489             if (bIfReset)
1490             {
1491                 for (i = 0; i < nlim; i++)
1492                 {
1493                     dfhist->wl_histo[i] = 0;
1494                 }
1495                 dfhist->wl_delta *= expand->wl_scale;
1496                 if (log)
1497                 {
1498                     fprintf(log, "\nStep %d: weights are now:", static_cast<int>(step));
1499                     for (i = 0; i < nlim; i++)
1500                     {
1501                         fprintf(log, " %.5f", dfhist->sum_weights[i]);
1502                     }
1503                     fprintf(log, "\n");
1504                 }
1505             }
1506         }
1507     }
1508     sfree(pfep_lamee);
1509     sfree(scaled_lamee);
1510     sfree(weighted_lamee);
1511     sfree(p_k);
1512
1513     return lamnew;
1514 }
1515
1516 //! Update reference temperature for simulated tempering state change
1517 static void simulatedTemperingUpdateTemperature(t_inputrec*                         ir,
1518                                                 t_state*                            state,
1519                                                 t_extmass*                          MassQ,
1520                                                 rvec*                               v,
1521                                                 const int                           homenr,
1522                                                 gmx::ArrayRef<const unsigned short> cTC,
1523                                                 const int                           lamnew)
1524 {
1525     const t_simtemp*  simtemp = ir->simtempvals.get();
1526     std::vector<real> buf_ngtc(ir->opts.ngtc);
1527
1528     for (int i = 0; i < ir->opts.ngtc; i++)
1529     {
1530         if (ir->opts.ref_t[i] > 0)
1531         {
1532             real told         = ir->opts.ref_t[i];
1533             ir->opts.ref_t[i] = simtemp->temperatures[lamnew];
1534             buf_ngtc[i] = std::sqrt(ir->opts.ref_t[i] / told); /* using the buffer as temperature scaling */
1535         }
1536     }
1537
1538     /* we don't need to manipulate the ekind information, as it isn't due to be reset until the next step anyway */
1539
1540     for (int n = 0; n < homenr; n++)
1541     {
1542         const int gt = cTC.empty() ? 0 : cTC[n];
1543         for (int d = 0; d < DIM; d++)
1544         {
1545             v[n][d] *= buf_ngtc[gt];
1546         }
1547     }
1548
1549     if (inputrecNptTrotter(ir) || inputrecNphTrotter(ir) || inputrecNvtTrotter(ir))
1550     {
1551         /* we need to recalculate the masses if the temperature has changed */
1552         init_npt_masses(ir, state, MassQ, FALSE);
1553         for (int i = 0; i < state->nnhpres; i++)
1554         {
1555             for (int j = 0; j < ir->opts.nhchainlength; j++)
1556             {
1557                 state->nhpres_vxi[i + j] *= buf_ngtc[i];
1558             }
1559         }
1560         for (int i = 0; i < ir->opts.ngtc; i++)
1561         {
1562             for (int j = 0; j < ir->opts.nhchainlength; j++)
1563             {
1564                 state->nosehoover_vxi[i + j] *= buf_ngtc[i];
1565             }
1566         }
1567     }
1568 }
1569
1570 int ExpandedEnsembleDynamics(FILE*                               log,
1571                              t_inputrec*                         ir,
1572                              const gmx_enerdata_t*               enerd,
1573                              t_state*                            state,
1574                              t_extmass*                          MassQ,
1575                              int                                 fep_state,
1576                              df_history_t*                       dfhist,
1577                              int64_t                             step,
1578                              rvec*                               v,
1579                              const int                           homenr,
1580                              gmx::ArrayRef<const unsigned short> cTC)
1581 /* Note that the state variable is only needed for simulated tempering, not
1582    Hamiltonian expanded ensemble.  May be able to remove it after integrator refactoring. */
1583 {
1584     const int newLambda = expandedEnsembleUpdateLambdaState(log, ir, enerd, fep_state, dfhist, step);
1585     // if using simulated tempering, we need to adjust the temperatures
1586     // only need to change the temperatures if we change the state
1587     if (ir->bSimTemp && (newLambda != fep_state))
1588     {
1589         simulatedTemperingUpdateTemperature(ir, state, MassQ, v, homenr, cTC, newLambda);
1590     }
1591     return newLambda;
1592 }