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