src/gromacs/applied_forces/awh/biaswriter.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2015,2016,2017,2018,2019,2020,2021, 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/applied_forces/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 "biasgrid.h"
  52 #include "correlationgrid.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     { AwhOutputEntryType::MetaData, Normalization::None },
  70     { AwhOutputEntryType::CoordValue, Normalization::Coordinate },
  71     { AwhOutputEntryType::Pmf, Normalization::FreeEnergy },
  72     { AwhOutputEntryType::Bias, Normalization::FreeEnergy },
  73     { AwhOutputEntryType::Visits, Normalization::Distribution },
  74     { AwhOutputEntryType::Weights, Normalization::Distribution },
  75     { AwhOutputEntryType::Target, Normalization::Distribution },
  76     { AwhOutputEntryType::ForceCorrelationVolume, Normalization::Distribution },
  77     { AwhOutputEntryType::FrictionTensor, Normalization::None }
  78 };
  79
  80 /*! \brief
  81  * Gets the coordinate normalization value for the given dimension.
  82  *
  83  * \param[in] bias      The AWH bias.
  84  * \param[in] dimIndex  Dimensional index.
  85  * \returns the coordinate normalization value.
  86  */
  87 float getCoordNormalizationValue(const Bias& bias, int dimIndex)
  88 {
  89     /* AWH may use different units internally but here we convert to user units */
  90     return bias.dimParams()[dimIndex].scaleInternalToUserInput(1);
  91 }
  92
  93 /*! \brief
  94  * Gets the normalization value for the given output entry type.
  95  *
  96  * \param[in] outputType  Output entry type.
  97  * \param[in] bias        The AWH bias.
  98  * \param[in] numBlocks   The number of blocks for this output type.
  99  * \returns the normalization value.
 100  */
 101 float getNormalizationValue(AwhOutputEntryType outputType, const Bias& bias, int numBlocks)
 102 {
 103     float normalizationValue = 0;
 104
 105     switch (outputType)
 106     {
 107         case AwhOutputEntryType::CoordValue:
 108             normalizationValue = getCoordNormalizationValue(bias, numBlocks);
 109             break;
 110         case AwhOutputEntryType::Visits:
 111         case AwhOutputEntryType::Weights:
 112         case AwhOutputEntryType::Target:
 113             normalizationValue = static_cast<float>(bias.state().points().size());
 114             break;
 115         case AwhOutputEntryType::ForceCorrelationVolume:
 116             normalizationValue = static_cast<double>(bias.state().points().size());
 117             break;
 118         default: break;
 119     }
 120
 121     return normalizationValue;
 122 }
 123
 124 } // namespace
 125
 126 AwhEnergyBlock::AwhEnergyBlock(int numPoints, Normalization normalizationType, float normalizationValue) :
 127     normalizationType(normalizationType), normalizationValue(normalizationValue), data_(numPoints)
 128 {
 129 }
 130
 131 BiasWriter::BiasWriter(const Bias& bias)
 132 {
 133     std::map<AwhOutputEntryType, int> outputTypeNumBlock; /* Number of blocks per output type */
 134
 135     /* Different output variable types need different number of blocks.
 136      * We keep track of the starting block for each variable.
 137      */
 138     int blockCount = 0;
 139     for (const auto& pair : outputTypeToNormalization)
 140     {
 141         const AwhOutputEntryType outputType = pair.first;
 142         {
 143             outputTypeToBlock_[outputType] = blockCount;
 144
 145             if (outputType == AwhOutputEntryType::CoordValue)
 146             {
 147                 outputTypeNumBlock[outputType] = bias.ndim();
 148             }
 149             else if (outputType == AwhOutputEntryType::FrictionTensor)
 150             {
 151                 outputTypeNumBlock[outputType] = bias.forceCorrelationGrid().tensorSize();
 152             }
 153             else
 154             {
 155                 /* Most output variable types need one block */
 156                 outputTypeNumBlock[outputType] = 1;
 157             }
 158         }
 159         blockCount += outputTypeNumBlock[outputType];
 160     }
 161
 162     /* Initialize the data blocks for each variable */
 163     for (const auto& pair : outputTypeToNormalization)
 164     {
 165         const AwhOutputEntryType outputType = pair.first;
 166         int                      numPoints;
 167         if (outputType == AwhOutputEntryType::MetaData)
 168         {
 169             numPoints = static_cast<int>(AwhOutputMetaData::Count);
 170         }
 171         else
 172         {
 173             numPoints = bias.state().points().size();
 174         }
 175         for (int b = 0; b < outputTypeNumBlock[outputType]; b++)
 176         {
 177             block_.emplace_back(numPoints, pair.second, getNormalizationValue(outputType, bias, b));
 178         }
 179     }
 180 }
 181
 182 /*! \brief
 183  * Normalizes block data for output.
 184  *
 185  * \param[in,out] block  The block to normalize.
 186  * \param[in]     bias   The AWH bias.
 187  */
 188 static void normalizeBlock(AwhEnergyBlock* block, const Bias& bias)
 189 {
 190     gmx::ArrayRef<float> data = block->data();
 191
 192     /* Here we operate on float data (which is accurate enough, since it
 193      * is statistical data that will never reach full float precision).
 194      * But since we can have very many data points, we sum into a double.
 195      */
 196     double sum       = 0;
 197     float  minValue  = GMX_FLOAT_MAX;
 198     float  recipNorm = 0;
 199
 200     switch (block->normalizationType)
 201     {
 202         case Normalization::None: break;
 203         case Normalization::Coordinate:
 204             /* Normalize coordinate values by a scale factor */
 205             for (float& point : data)
 206             {
 207                 point *= block->normalizationValue;
 208             }
 209             break;
 210         case Normalization::FreeEnergy:
 211             /* Normalize free energy values by subtracting the minimum value */
 212             for (gmx::index index = 0; index < data.ssize(); index++)
 213             {
 214                 if (bias.state().points()[index].inTargetRegion() && data[index] < minValue)
 215                 {
 216                     minValue = data[index];
 217                 }
 218             }
 219             for (gmx::index index = 0; index < data.ssize(); index++)
 220             {
 221                 if (bias.state().points()[index].inTargetRegion())
 222                 {
 223                     data[index] -= minValue;
 224                 }
 225             }
 226
 227             break;
 228         case Normalization::Distribution:
 229             /* Normalize distribution values by normalizing their sum */
 230             for (float& point : data)
 231             {
 232                 sum += point;
 233             }
 234             if (sum > 0)
 235             {
 236                 recipNorm = block->normalizationValue / static_cast<float>(sum);
 237             }
 238             for (float& point : data)
 239             {
 240                 point *= recipNorm;
 241             }
 242             break;
 243         default: GMX_RELEASE_ASSERT(false, "Unknown AWH normalization type"); break;
 244     }
 245 }
 246
 247 void BiasWriter::transferMetaDataToWriter(gmx::index        metaDataIndex,
 248                                           AwhOutputMetaData metaDataType,
 249                                           const Bias&       bias)
 250 {
 251     gmx::ArrayRef<float> data = block_[getVarStartBlock(AwhOutputEntryType::MetaData)].data();
 252     GMX_ASSERT(metaDataIndex < data.ssize(),
 253                "Attempt to transfer AWH meta data to block for index out of range");
 254
 255     /* Transfer the point data of this variable to the right block(s) */
 256     switch (metaDataType)
 257     {
 258         case AwhOutputMetaData::NumBlock:
 259             /* The number of subblocks per awh (needed by gmx_energy) */
 260             data[metaDataIndex] = static_cast<double>(block_.size());
 261             /* Note: a single subblock takes only a single type and we need doubles. */
 262             break;
 263         case AwhOutputMetaData::TargetError:
 264             /* The theoretical target error */
 265             data[metaDataIndex] = bias.params().initialErrorInKT
 266                                   * std::sqrt(bias.params().initialHistogramSize
 267                                               / bias.state().histogramSize().histogramSize());
 268             break;
 269         case AwhOutputMetaData::ScaledSampleWeight:
 270             /* The logarithm of the sample weight relative to a sample weight of 1 at the initial time.
 271                In the normal case: this will increase in the initial stage and then stay at a constant value. */
 272             data[metaDataIndex] = bias.state().histogramSize().logScaledSampleWeight();
 273             break;
 274         case AwhOutputMetaData::Count: break;
 275     }
 276 }
 277
 278 void BiasWriter::transferPointDataToWriter(AwhOutputEntryType         outputType,
 279                                            int                        pointIndex,
 280                                            const Bias&                bias,
 281                                            gmx::ArrayRef<const float> pmf)
 282 {
 283     /* The starting block index of this output type.
 284      * Note that some variables need several (contiguous) blocks.
 285      */
 286     int blockStart = getVarStartBlock(outputType);
 287     GMX_ASSERT(pointIndex < static_cast<int>(block_[blockStart].data().size()),
 288                "Attempt to transfer AWH data to block for point index out of range");
 289
 290     const CorrelationGrid& forceCorrelation = bias.forceCorrelationGrid();
 291     int                    numCorrelation   = forceCorrelation.tensorSize();
 292
 293     /* Transfer the point data of this variable to the right block(s) */
 294     int b = blockStart;
 295     switch (outputType)
 296     {
 297         case AwhOutputEntryType::MetaData:
 298             GMX_RELEASE_ASSERT(false, "MetaData is handled by a different function");
 299             break;
 300         case AwhOutputEntryType::CoordValue:
 301         {
 302             const awh_dvec& coordValue = bias.getGridCoordValue(pointIndex);
 303             for (int d = 0; d < bias.ndim(); d++)
 304             {
 305                 block_[b].data()[pointIndex] = coordValue[d];
 306                 b++;
 307             }
 308         }
 309         break;
 310         case AwhOutputEntryType::Pmf:
 311             block_[b].data()[pointIndex] =
 312                     bias.state().points()[pointIndex].inTargetRegion() ? pmf[pointIndex] : 0;
 313             break;
 314         case AwhOutputEntryType::Bias:
 315         {
 316             const awh_dvec& coordValue   = bias.getGridCoordValue(pointIndex);
 317             block_[b].data()[pointIndex] = bias.state().points()[pointIndex].inTargetRegion()
 318                                                    ? bias.calcConvolvedBias(coordValue)
 319                                                    : 0;
 320         }
 321         break;
 322         case AwhOutputEntryType::Visits:
 323             block_[b].data()[pointIndex] = bias.state().points()[pointIndex].numVisitsTot();
 324             break;
 325         case AwhOutputEntryType::Weights:
 326             block_[b].data()[pointIndex] = bias.state().points()[pointIndex].weightSumTot();
 327             break;
 328         case AwhOutputEntryType::Target:
 329             block_[b].data()[pointIndex] = bias.state().points()[pointIndex].target();
 330             break;
 331         case AwhOutputEntryType::ForceCorrelationVolume:
 332             block_[b].data()[pointIndex] =
 333                     forceCorrelation.tensors()[pointIndex].getVolumeElement(forceCorrelation.dtSample);
 334             break;
 335         case AwhOutputEntryType::FrictionTensor:
 336             /* Store force correlation in units of friction, i.e. time/length^2 */
 337             for (int n = 0; n < numCorrelation; n++)
 338             {
 339                 block_[b].data()[pointIndex] = forceCorrelation.tensors()[pointIndex].getTimeIntegral(
 340                         n, forceCorrelation.dtSample);
 341                 b++;
 342             }
 343             break;
 344         default: GMX_RELEASE_ASSERT(false, "Unknown AWH output variable"); break;
 345     }
 346 }
 347
 348 void BiasWriter::prepareBiasOutput(const Bias& bias)
 349 {
 350     /* Pack the AWH data into the writer data. */
 351
 352     /* Evaluate the PMF for all points */
 353     gmx::ArrayRef<float> pmf = block_[getVarStartBlock(AwhOutputEntryType::Pmf)].data();
 354     bias.state().getPmf(pmf);
 355
 356     /* Pack the data point by point.
 357      * Unfortunately we can not loop over a class enum, so we cast to int.
 358      * \todo Use strings instead of enum when we port the output to TNG.
 359      */
 360     for (int i = 0; i < static_cast<int>(AwhOutputMetaData::Count); i++)
 361     {
 362         transferMetaDataToWriter(i, static_cast<AwhOutputMetaData>(i), bias);
 363     }
 364     for (const auto& pair : outputTypeToNormalization)
 365     {
 366         const AwhOutputEntryType outputType = pair.first;
 367         /* Skip metadata (transfered above) and unused blocks */
 368         if (outputType == AwhOutputEntryType::MetaData || !hasVarBlock(outputType))
 369         {
 370             continue;
 371         }
 372         for (size_t m = 0; m < bias.state().points().size(); m++)
 373         {
 374             transferPointDataToWriter(outputType, m, bias, pmf);
 375         }
 376     }
 377
 378     /* For looks of the output, normalize it */
 379     for (AwhEnergyBlock& block : block_)
 380     {
 381         normalizeBlock(&block, bias);
 382     }
 383 }
 384
 385 int BiasWriter::writeToEnergySubblocks(const Bias& bias, t_enxsubblock* sub)
 386 {
 387     prepareBiasOutput(bias);
 388
 389     for (size_t b = 0; b < block_.size(); b++)
 390     {
 391         sub[b].type = XdrDataType::Float;
 392         sub[b].nr   = block_[b].data().size();
 393         sub[b].fval = block_[b].data().data();
 394     }
 395
 396     return block_.size();
 397 }
 398
 399 } // namespace gmx