src/gromacs/awh/biaswriter.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2015,2016,2017,2018, by the GROMACS development team, led by
   5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6  * and including many others, as listed in the AUTHORS file in the
   7  * top-level source directory and at http://www.gromacs.org.
   8  *
   9  * GROMACS is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public License
  11  * as published by the Free Software Foundation; either version 2.1
  12  * of the License, or (at your option) any later version.
  13  *
  14  * GROMACS is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with GROMACS; if not, see
  21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23  *
  24  * If you want to redistribute modifications to GROMACS, please
  25  * consider that scientific software is very special. Version
  26  * control is crucial - bugs must be traceable. We will be happy to
  27  * consider code for inclusion in the official distribution, but
  28  * derived work must not be called official GROMACS. Details are found
  29  * in the README & COPYING files - if they are missing, get the
  30  * official version at http://www.gromacs.org.
  31  *
  32  * To help us fund GROMACS development, we humbly ask that you cite
  33  * the research papers on the package. Check out http://www.gromacs.org.
  34  */
  35
  36 #include "gmxpre.h"
  37
  38 #include "biaswriter.h"
  39
  40 #include <cassert>
  41 #include <cmath>
  42
  43 #include "gromacs/awh/awh.h"
  44 #include "gromacs/mdtypes/awh-params.h"
  45 #include "gromacs/mdtypes/commrec.h"
  46 #include "gromacs/trajectory/energyframe.h"
  47 #include "gromacs/utility/gmxassert.h"
  48 #include "gromacs/utility/smalloc.h"
  49
  50 #include "bias.h"
  51 #include "correlationgrid.h"
  52 #include "grid.h"
  53 #include "pointstate.h"
  54
  55 namespace gmx
  56 {
  57
  58 namespace
  59 {
  60
  61 /*! \brief
  62  * Map the output entry type to a normalization type.
  63  *
  64  * The data is written to energy file blocks in the order given by
  65  * the iterator of this map, which is based on the enum value
  66  * (and matches the order of the lines below).
  67  */
  68 const std::map<AwhOutputEntryType, Normalization> outputTypeToNormalization =
  69 {
  70     { AwhOutputEntryType::MetaData,               Normalization::None },
  71     { AwhOutputEntryType::CoordValue,             Normalization::Coordinate },
  72     { AwhOutputEntryType::Pmf,                    Normalization::FreeEnergy },
  73     { AwhOutputEntryType::Bias,                   Normalization::FreeEnergy },
  74     { AwhOutputEntryType::Visits,                 Normalization::Distribution },
  75     { AwhOutputEntryType::Weights,                Normalization::Distribution },
  76     { AwhOutputEntryType::Target,                 Normalization::Distribution },
  77     { AwhOutputEntryType::ForceCorrelationVolume, Normalization::Distribution },
  78     { AwhOutputEntryType::FrictionTensor,         Normalization::None }
  79 };
  80
  81 /*! \brief
  82  * Gets the coordinate normalization value for the given dimension.
  83  *
  84  * \param[in] bias      The AWH bias.
  85  * \param[in] dimIndex  Dimensional index.
  86  * \returns the coordinate normalization value.
  87  */
  88 float getCoordNormalizationValue(const Bias &bias,
  89                                  int         dimIndex)
  90 {
  91     /* AWH may use different units internally but here we convert to user units */
  92     return bias.dimParams()[dimIndex].scaleInternalToUserInput(1);
  93 }
  94
  95 /*! \brief
  96  * Gets the normalization value for the given output entry type.
  97  *
  98  * \param[in] outputType  Output entry type.
  99  * \param[in] bias        The AWH bias.
 100  * \param[in] numBlocks   The number of blocks for this output type.
 101  * \returns the normalization value.
 102  */
 103 float getNormalizationValue(AwhOutputEntryType  outputType,
 104                             const Bias         &bias,
 105                             int                 numBlocks)
 106 {
 107     float normalizationValue = 0;
 108
 109     switch (outputType)
 110     {
 111         case AwhOutputEntryType::CoordValue:
 112             normalizationValue = getCoordNormalizationValue(bias, numBlocks);
 113             break;
 114         case AwhOutputEntryType::Visits:
 115         case AwhOutputEntryType::Weights:
 116         case AwhOutputEntryType::Target:
 117             normalizationValue = static_cast<float>(bias.state().points().size());
 118             break;
 119         case AwhOutputEntryType::ForceCorrelationVolume:
 120             normalizationValue = static_cast<double>(bias.state().points().size());
 121             break;
 122         default:
 123             break;
 124     }
 125
 126     return normalizationValue;
 127 }
 128
 129 }   // namespace
 130
 131 AwhEnergyBlock::AwhEnergyBlock(int            numPoints,
 132                                Normalization  normalizationType,
 133                                float          normalizationValue) :
 134     normalizationType(normalizationType),
 135     normalizationValue(normalizationValue),
 136     data_(numPoints)
 137 {
 138 }
 139
 140 BiasWriter::BiasWriter(const Bias &bias)
 141 {
 142     std::map<AwhOutputEntryType, int> outputTypeNumBlock; /* Number of blocks per output type */
 143
 144     /* Different output variable types need different number of blocks.
 145      * We keep track of the starting block for each variable.
 146      */
 147     int blockCount = 0;
 148     for (const auto &pair : outputTypeToNormalization)
 149     {
 150         const AwhOutputEntryType outputType = pair.first;
 151         {
 152             outputTypeToBlock_[outputType] = blockCount;
 153
 154             if (outputType == AwhOutputEntryType::CoordValue)
 155             {
 156                 outputTypeNumBlock[outputType] = bias.ndim();
 157             }
 158             else if (outputType == AwhOutputEntryType::FrictionTensor)
 159             {
 160                 outputTypeNumBlock[outputType] = bias.forceCorrelationGrid().tensorSize();
 161             }
 162             else
 163             {
 164                 /* Most output variable types need one block */
 165                 outputTypeNumBlock[outputType] = 1;
 166             }
 167         }
 168         blockCount += outputTypeNumBlock[outputType];
 169     }
 170
 171     /* Initialize the data blocks for each variable */
 172     for (const auto &pair : outputTypeToNormalization)
 173     {
 174         const AwhOutputEntryType outputType = pair.first;
 175         int                      numPoints;
 176         if (outputType == AwhOutputEntryType::MetaData)
 177         {
 178             numPoints = static_cast<int>(AwhOutputMetaData::Count);
 179         }
 180         else
 181         {
 182             numPoints = bias.state().points().size();
 183         }
 184         for (int b = 0; b < outputTypeNumBlock[outputType]; b++)
 185         {
 186             block_.emplace_back(numPoints,
 187                                 pair.second,
 188                                 getNormalizationValue(outputType, bias, b));
 189         }
 190     }
 191 }
 192
 193 /*! \brief
 194  * Normalizes block data for output.
 195  *
 196  * \param[in,out] block  The block to normalize.
 197  * \param[in]     bias   The AWH bias.
 198  */
 199 static void normalizeBlock(AwhEnergyBlock *block, const Bias &bias)
 200 {
 201     gmx::ArrayRef<float> data = block->data();
 202
 203     /* Here we operate on float data (which is accurate enough, since it
 204      * is statistical data that will never reach full float precision).
 205      * But since we can have very many data points, we sum into a double.
 206      */
 207     double sum       = 0;
 208     float  minValue  = GMX_FLOAT_MAX;
 209     float  recipNorm = 0;
 210
 211     switch (block->normalizationType)
 212     {
 213         case Normalization::None:
 214             break;
 215         case Normalization::Coordinate:
 216             /* Normalize coordinate values by a scale factor */
 217             for (float &point : data)
 218             {
 219                 point *= block->normalizationValue;
 220             }
 221             break;
 222         case Normalization::FreeEnergy:
 223             /* Normalize free energy values by subtracting the minimum value */
 224             for (gmx::index index = 0; index < data.size(); index++)
 225             {
 226                 if (bias.state().points()[index].inTargetRegion() && data[index] < minValue)
 227                 {
 228                     minValue = data[index];
 229                 }
 230             }
 231             for (gmx::index index = 0; index < data.size(); index++)
 232             {
 233                 if (bias.state().points()[index].inTargetRegion())
 234                 {
 235                     data[index] -= minValue;
 236                 }
 237             }
 238
 239             break;
 240         case Normalization::Distribution:
 241             /* Normalize distribution values by normalizing their sum */
 242             for (float &point : data)
 243             {
 244                 sum += point;
 245             }
 246             if (sum > 0)
 247             {
 248                 recipNorm = block->normalizationValue/static_cast<float>(sum);
 249             }
 250             for (float &point : data)
 251             {
 252                 point *= recipNorm;
 253             }
 254             break;
 255         default:
 256             GMX_RELEASE_ASSERT(false, "Unknown AWH normalization type");
 257             break;
 258     }
 259 }
 260
 261 void BiasWriter::transferMetaDataToWriter(gmx::index         metaDataIndex,
 262                                           AwhOutputMetaData  metaDataType,
 263                                           const Bias        &bias)
 264 {
 265     gmx::ArrayRef<float> data = block_[getVarStartBlock(AwhOutputEntryType::MetaData)].data();
 266     GMX_ASSERT(metaDataIndex < data.size(), "Attempt to transfer AWH meta data to block for index out of range");
 267
 268     /* Transfer the point data of this variable to the right block(s) */
 269     switch (metaDataType)
 270     {
 271         case AwhOutputMetaData::NumBlock:
 272             /* The number of subblocks per awh (needed by gmx_energy) */
 273             data[metaDataIndex] = static_cast<double>(block_.size());
 274             /* Note: a single subblock takes only a single type and we need doubles. */
 275             break;
 276         case AwhOutputMetaData::TargetError:
 277             /* The theoretical target error */
 278             data[metaDataIndex] = bias.params().initialErrorInKT*std::sqrt(bias.params().initialHistogramSize/bias.state().histogramSize().histogramSize());
 279             break;
 280         case AwhOutputMetaData::ScaledSampleWeight:
 281             /* The logarithm of the sample weight relative to a sample weight of 1 at the initial time.
 282                In the normal case: this will increase in the initial stage and then stay at a constant value. */
 283             data[metaDataIndex] = bias.state().histogramSize().logScaledSampleWeight();
 284             break;
 285         case AwhOutputMetaData::Count:
 286             break;
 287     }
 288 }
 289
 290 void
 291 BiasWriter::transferPointDataToWriter(AwhOutputEntryType          outputType,
 292                                       int                         pointIndex,
 293                                       const Bias                 &bias,
 294                                       gmx::ArrayRef<const float>  pmf)
 295 {
 296     /* The starting block index of this output type.
 297      * Note that some variables need several (contiguous) blocks.
 298      */
 299     int blockStart = getVarStartBlock(outputType);
 300     GMX_ASSERT(pointIndex < static_cast<int>(block_[blockStart].data().size()), "Attempt to transfer AWH data to block for point index out of range");
 301
 302     const CorrelationGrid &forceCorrelation = bias.forceCorrelationGrid();
 303     int                    numCorrelation   = forceCorrelation.tensorSize();
 304
 305     /* Transfer the point data of this variable to the right block(s) */
 306     int b = blockStart;
 307     switch (outputType)
 308     {
 309         case AwhOutputEntryType::MetaData:
 310             GMX_RELEASE_ASSERT(false, "MetaData is handled by a different function");
 311             break;
 312         case AwhOutputEntryType::CoordValue:
 313         {
 314             const awh_dvec &coordValue = bias.getGridCoordValue(pointIndex);
 315             for (int d = 0; d < bias.ndim(); d++)
 316             {
 317                 block_[b].data()[pointIndex] = coordValue[d];
 318                 b++;
 319             }
 320         }
 321         break;
 322         case AwhOutputEntryType::Pmf:
 323             block_[b].data()[pointIndex] = bias.state().points()[pointIndex].inTargetRegion() ? pmf[pointIndex] : 0;
 324             break;
 325         case AwhOutputEntryType::Bias:
 326         {
 327             const awh_dvec &coordValue = bias.getGridCoordValue(pointIndex);
 328             block_[b].data()[pointIndex] = bias.state().points()[pointIndex].inTargetRegion() ? bias.calcConvolvedBias(coordValue) : 0;
 329         }
 330         break;
 331         case AwhOutputEntryType::Visits:
 332             block_[b].data()[pointIndex] = bias.state().points()[pointIndex].numVisitsTot();
 333             break;
 334         case AwhOutputEntryType::Weights:
 335             block_[b].data()[pointIndex] = bias.state().points()[pointIndex].weightSumTot();
 336             break;
 337         case AwhOutputEntryType::Target:
 338             block_[b].data()[pointIndex] = bias.state().points()[pointIndex].target();
 339             break;
 340         case AwhOutputEntryType::ForceCorrelationVolume:
 341             block_[b].data()[pointIndex] = forceCorrelation.tensors()[pointIndex].getVolumeElement(forceCorrelation.dtSample);
 342             break;
 343         case AwhOutputEntryType::FrictionTensor:
 344             /* Store force correlation in units of friction, i.e. time/length^2 */
 345             for (int n = 0; n < numCorrelation; n++)
 346             {
 347                 block_[b].data()[pointIndex] = forceCorrelation.tensors()[pointIndex].getTimeIntegral(n, forceCorrelation.dtSample);
 348                 b++;
 349             }
 350             break;
 351         default:
 352             GMX_RELEASE_ASSERT(false, "Unknown AWH output variable");
 353             break;
 354     }
 355 }
 356
 357 void BiasWriter::prepareBiasOutput(const Bias &bias)
 358 {
 359     /* Pack the AWH data into the writer data. */
 360
 361     /* Evaluate the PMF for all points */
 362     gmx::ArrayRef<float> pmf = block_[getVarStartBlock(AwhOutputEntryType::Pmf)].data();
 363     bias.state().getPmf(pmf);
 364
 365     /* Pack the the data point by point.
 366      * Unfortunately we can not loop over a class enum, so we cast to int.
 367      * \todo Use strings instead of enum when we port the output to TNG.
 368      */
 369     for (int i = 0; i < static_cast<int>(AwhOutputMetaData::Count); i++)
 370     {
 371         transferMetaDataToWriter(i, static_cast<AwhOutputMetaData>(i), bias);
 372     }
 373     for (const auto &pair : outputTypeToNormalization)
 374     {
 375         const AwhOutputEntryType outputType = pair.first;
 376         /* Skip metadata (transfered above) and unused blocks */
 377         if (outputType == AwhOutputEntryType::MetaData || !hasVarBlock(outputType))
 378         {
 379             continue;
 380         }
 381         for (size_t m = 0; m < bias.state().points().size(); m++)
 382         {
 383             transferPointDataToWriter(outputType, m, bias, pmf);
 384         }
 385     }
 386
 387     /* For looks of the output, normalize it */
 388     for (AwhEnergyBlock &block : block_)
 389     {
 390         normalizeBlock(&block, bias);
 391     }
 392 }
 393
 394 int BiasWriter::writeToEnergySubblocks(const Bias    &bias,
 395                                        t_enxsubblock *sub)
 396 {
 397     prepareBiasOutput(bias);
 398
 399     for (size_t b = 0; b < block_.size(); b++)
 400     {
 401         sub[b].type = xdr_datatype_float;
 402         sub[b].nr   = block_[b].data().size();
 403         sub[b].fval = block_[b].data().data();
 404     }
 405
 406     return block_.size();
 407 }
 408
 409 }  // namespace gmx