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