Make it possible to use FEP lambda states as a reaction coordinate in AWH. Atom masse...

[alexxy/gromacs.git] / src / gromacs / awh / bias.h

diff --git a/src/gromacs/awh/bias.h b/src/gromacs/awh/bias.h
index de16d7a7c66aa7683c1b0f14a62c7464979d2fc0..d034aea1cf0da6140f3e3d4dc61d5e2bf5162cd7 100644 (file)
--- a/src/gromacs/awh/bias.h
+++ b/src/gromacs/awh/bias.h
@@ -159,7 +159,8 @@ public:
      * \param[in] mdTimeStep             The MD time step.
      * \param[in] numSharingSimulations  The number of simulations to share the bias across.
      * \param[in] biasInitFilename       Name of file to read PMF and target from.
-     * \param[in] thisRankWillDoIO       Tells whether this MPI rank will do I/O (checkpointing, AWH output), normally (only) the master rank does I/O.
+     * \param[in] thisRankWillDoIO       Tells whether this MPI rank will do I/O (checkpointing, AWH output),
+     * normally (only) the master rank does I/O.
      * \param[in] disableUpdateSkips     If to disable update skips, useful for testing.
      */
     Bias(int                            biasIndexInCollection,
@@ -192,6 +193,18 @@ public:
      * - reweight samples to extract the PMF.
      *
      * \param[in]     coordValue     The current coordinate value(s).
+     * \param[in]     neighborLambdaEnergies An array containing the energy of the system
+     * in neighboring lambdas. The array is of length numLambdas+1, where numLambdas is
+     * the number of free energy lambda states. Element 0 in the array is the energy
+     * of the current state and elements 1..numLambdas contain the energy of the system in the
+     * neighboring lambda states (also including the current state). When there are no free
+     * energy lambda state dimensions this can be empty.
+     * \param[in]     neighborLambdaDhdl     An array containing the dHdL at the neighboring lambda
+     * points. The array is of length numLambdas+1, where numLambdas is the number of free
+     * energy lambda states. Element 0 in the array is the dHdL
+     * of the current state and elements 1..numLambdas contain the dHdL of the system in the
+     * neighboring lambda states (also including the current state). When there are no free
+     * energy lambda state dimensions this can be empty.
      * \param[out]    awhPotential   Bias potential.
      * \param[out]    potentialJump  Change in bias potential for this bias.
      * \param[in]     commRecord     Struct for intra-simulation communication.
@@ -202,15 +215,17 @@ public:
      * \param[in,out] fplog          Log file.
      * \returns a reference to the bias force, size \ref ndim(), valid until the next call of this method or destruction of Bias, whichever comes first.
      */
-    gmx::ArrayRef<const double> calcForceAndUpdateBias(const awh_dvec        coordValue,
-                                                       double*               awhPotential,
-                                                       double*               potentialJump,
-                                                       const t_commrec*      commRecord,
-                                                       const gmx_multisim_t* ms,
-                                                       double                t,
-                                                       int64_t               step,
-                                                       int64_t               seed,
-                                                       FILE*                 fplog);
+    gmx::ArrayRef<const double> calcForceAndUpdateBias(const awh_dvec         coordValue,
+                                                       ArrayRef<const double> neighborLambdaEnergies,
+                                                       ArrayRef<const double> neighborLambdaDhdl,
+                                                       double*                awhPotential,
+                                                       double*                potentialJump,
+                                                       const t_commrec*       commRecord,
+                                                       const gmx_multisim_t*  ms,
+                                                       double                 t,
+                                                       int64_t                step,
+                                                       int64_t                seed,
+                                                       FILE*                  fplog);
 
     /*! \brief
      * Calculates the convolved bias for a given coordinate value.
@@ -301,9 +316,17 @@ private:
      * Collect samples for the force correlation analysis on the grid.
      *
      * \param[in] probWeightNeighbor  Probability weight of the neighboring points.
+     * \param[in] neighborLambdaDhdl  An array containing the dHdL at the neighboring lambda
+     * points. The array is of length numLambdas+1, where numLambdas is the number of free
+     * energy lambda states. Element 0 in the array is the dHdL
+     * of the current state and elements 1..numLambdas contain the dHdL of the system in the
+     * neighboring lambda states (also including the current state). When there are no free
+     * energy lambda state dimensions this can be empty.
      * \param[in] t                   The time.
      */
-    void updateForceCorrelationGrid(gmx::ArrayRef<const double> probWeightNeighbor, double t);
+    void updateForceCorrelationGrid(gmx::ArrayRef<const double> probWeightNeighbor,
+                                    ArrayRef<const double>      neighborLambdaDhdl,
+                                    double                      t);
 
 public:
     /*! \brief Return a const reference to the force correlation grid.
@@ -328,6 +351,30 @@ public:
      */
     int writeToEnergySubblocks(t_enxsubblock* subblock) const;
 
+    /*! \brief Returns true if the bias has a free energy lambda state dimension at all.
+     */
+    bool hasFepLambdaDimension() const
+    {
+        return std::any_of(std::begin(dimParams_), std::end(dimParams_),
+                           [](const auto& dimParam) { return dimParam.isFepLambdaDimension(); });
+    }
+
+    /*! \brief Returns whether the specified dimension is a free energy lambda
+     * state dimension.
+     *
+     * \param[in] dim      The dimension to check.
+     *
+     * \returns true if the specified dimension is a free energy lambda state dimension.
+     */
+    bool isFepLambdaDimension(int dim) const;
+
+    /*! \brief
+     * Returns whether we should sample the coordinate.
+     *
+     * \param[in] step  The MD step number.
+     */
+    bool isSampleCoordStep(int64_t step) const;
+
     /* Data members. */
 private:
     const std::vector<DimParams> dimParams_; /**< Parameters for each dimension. */