eec2648b28ee36f4c17d57e447e83e67afd18d1f
[alexxy/gromacs.git] / python_packaging / sample_restraint / src / cpp / ensemblepotential.h
1 #ifndef HARMONICRESTRAINT_ENSEMBLEPOTENTIAL_H
2 #define HARMONICRESTRAINT_ENSEMBLEPOTENTIAL_H
3
4 /*! \file
5  * \brief Provide restrained ensemble MD potential for GROMACS plugin.
6  *
7  * The restraint implemented here uses a facility provided by gmxapi to perform averaging of some
8  * array data across an ensemble of simulations. Simpler pair restraints can use less of this
9  * example code.
10  *
11  * Contains a lot of boiler plate that is being generalized and migrate out of this file, but other
12  * pair restraints can be implemented by following the example in this and ``ensemblepotential.cpp``.
13  * The ``CMakeLists.txt`` file will need to be updated if you add additional source files, and
14  * ``src/pythonmodule/export_plugin.cpp`` will need to be updated if you add or change the name of
15  * potentials.
16  *
17  * \author M. Eric Irrgang <ericirrgang@gmail.com>
18  */
19
20 #include <array>
21 #include <memory>
22 #include <mutex>
23 #include <vector>
24
25 #include "gmxapi/gromacsfwd.h"
26 #include "gmxapi/session.h"
27 #include "gmxapi/md/mdmodule.h"
28
29 #include "gromacs/restraint/restraintpotential.h"
30 #include "gromacs/utility/real.h"
31
32 #include "sessionresources.h"
33
34 namespace plugin
35 {
36
37 // Histogram for a single restrained pair.
38 using PairHist = std::vector<double>;
39
40 struct ensemble_input_param_type
41 {
42     /// distance histogram parameters
43     size_t nBins{0};
44     double binWidth{0.};
45
46     /// Flat-bottom potential boundaries.
47     double minDist{0};
48     double maxDist{0};
49
50     /// Experimental reference distribution.
51     PairHist experimental{};
52
53     /// Number of samples to store during each window.
54     unsigned int nSamples{0};
55     double samplePeriod{0};
56
57     /// Number of windows to use for smoothing histogram updates.
58     unsigned int nWindows{0};
59
60     /// Harmonic force coefficient
61     double k{0};
62     /// Smoothing factor: width of Gaussian interpolation for histogram
63     double sigma{0};
64
65 };
66
67 // \todo We should be able to automate a lot of the parameter setting stuff
68 // by having the developer specify a map of parameter names and the corresponding type, but that could get tricky.
69 // The statically compiled fast parameter structure would be generated with a recursive variadic template
70 // the way a tuple is. ref https://eli.thegreenplace.net/2014/variadic-templates-in-c/
71
72 std::unique_ptr<ensemble_input_param_type>
73 makeEnsembleParams(size_t nbins,
74                    double binWidth,
75                    double minDist,
76                    double maxDist,
77                    const std::vector<double>& experimental,
78                    unsigned int nSamples,
79                    double samplePeriod,
80                    unsigned int nWindows,
81                    double k,
82                    double sigma);
83
84 /*!
85  * \brief a residue-pair bias calculator for use in restrained-ensemble simulations.
86  *
87  * Applies a force between two sites according to the difference between an experimentally observed
88  * site pair distance distribution and the distance distribution observed earlier in the simulation
89  * trajectory. The sampled distribution is averaged from the previous `nwindows` histograms from all
90  * ensemble members. Each window contains a histogram populated with `nsamples` distances recorded at
91  * `sample_period` step intervals.
92  *
93  * \internal
94  * During a the window_update_period steps of a window, the potential applied is a harmonic function of
95  * the difference between the sampled and experimental histograms. At the beginning of the window, this
96  * difference is found and a Gaussian blur is applied.
97  */
98 class EnsemblePotential
99 {
100     public:
101         using input_param_type = ensemble_input_param_type;
102
103         /* No default constructor. Parameters must be provided. */
104         EnsemblePotential() = delete;
105
106         /*!
107          * \brief Constructor called by the wrapper code to produce a new instance.
108          *
109          * This constructor is called once per simulation per GROMACS process. Note that until
110          * gmxapi 0.0.8 there is only one instance per simulation in a thread-MPI simulation.
111          *
112          * \param params
113          */
114         explicit EnsemblePotential(const input_param_type& params);
115
116         /*!
117          * \brief Deprecated constructor taking a parameter list.
118          *
119          * \param nbins
120          * \param binWidth
121          * \param minDist
122          * \param maxDist
123          * \param experimental
124          * \param nSamples
125          * \param samplePeriod
126          * \param nWindows
127          * \param k
128          * \param sigma
129          */
130         EnsemblePotential(size_t nbins,
131                           double binWidth,
132                           double minDist,
133                           double maxDist,
134                           PairHist experimental,
135                           unsigned int nSamples,
136                           double samplePeriod,
137                           unsigned int nWindows,
138                           double k,
139                           double sigma);
140
141         /*!
142          * \brief Evaluates the pair restraint potential.
143          *
144          * In parallel simulations, the gmxapi framework does not make guarantees about where or
145          * how many times this function is called. It should be simple and stateless; it should not
146          * update class member data (see ``ensemblepotential.cpp``. For a more controlled API hook
147          * and to manage state in the object, use ``callback()``.
148          *
149          * \param v position of the site for which force is being calculated.
150          * \param v0 reference site (other member of the pair).
151          * \param t current simulation time (ps).
152          * \return container for force and potential energy data.
153          */
154         // Implementation note for the future: If dispatching this virtual function is not fast
155         // enough, the compiler may be able to better optimize a free
156         // function that receives the current restraint as an argument.
157         gmx::PotentialPointData calculate(gmx::Vector v,
158                                           gmx::Vector v0,
159                                           gmx_unused double t);
160
161         /*!
162          * \brief An update function to be called on the simulation master rank/thread periodically by the Restraint framework.
163          *
164          * Defining this function in a plugin potential is optional. If the function is defined,
165          * the restraint framework calls this function (on the first rank only in a parallel simulation) before calling calculate().
166          *
167          * The callback may use resources provided by the Session in the callback to perform updates
168          * to the local or global state of an ensemble of simulations. Future gmxapi releases will
169          * include additional optimizations, allowing call-back frequency to be expressed, and more
170          * general Session resources, as well as more flexible call signatures.
171          */
172         void callback(gmx::Vector v,
173                       gmx::Vector v0,
174                       double t,
175                       const Resources& resources);
176
177     private:
178         /// Width of bins (distance) in histogram
179         size_t nBins_;
180         double binWidth_;
181
182         /// Flat-bottom potential boundaries.
183         double minDist_;
184         double maxDist_;
185         /// Smoothed historic distribution for this restraint. An element of the array of restraints in this simulation.
186         // Was `hij` in earlier code.
187         PairHist histogram_;
188         PairHist experimental_;
189
190         /// Number of samples to store during each window.
191         unsigned int nSamples_;
192         unsigned int currentSample_;
193         double samplePeriod_;
194         double nextSampleTime_;
195         /// Accumulated list of samples during a new window.
196         std::vector<double> distanceSamples_;
197
198         /// Number of windows to use for smoothing histogram updates.
199         size_t nWindows_;
200         size_t currentWindow_;
201         double windowStartTime_;
202         double nextWindowUpdateTime_;
203         /// The history of nwindows histograms for this restraint.
204         std::vector<std::unique_ptr<plugin::Matrix<double>>> windows_;
205
206         /// Harmonic force coefficient
207         double k_;
208         /// Smoothing factor: width of Gaussian interpolation for histogram
209         double sigma_;
210 };
211
212 /*!
213  * \brief Use EnsemblePotential to implement a RestraintPotential
214  *
215  * This is boiler plate that will be templated and moved.
216  */
217 class EnsembleRestraint : public ::gmx::IRestraintPotential, private EnsemblePotential
218 {
219     public:
220         using EnsemblePotential::input_param_type;
221
222         EnsembleRestraint(std::vector<int> sites,
223                           const input_param_type& params,
224                           std::shared_ptr<Resources> resources
225         ) :
226             EnsemblePotential(params),
227             sites_{std::move(sites)},
228             resources_{std::move(resources)}
229         {}
230
231         ~EnsembleRestraint() override = default;
232
233         /*!
234          * \brief Implement required interface of gmx::IRestraintPotential
235          *
236          * \return list of configured site indices.
237          *
238          * \todo remove to template header
239          * \todo abstraction of site references
240          */
241         std::vector<int> sites() const override
242         {
243             return sites_;
244         }
245
246         /*!
247          * \brief Implement the interface gmx::IRestraintPotential
248          *
249          * Dispatch to calculate() method.
250          *
251          * \param r1 coordinate of first site
252          * \param r2 reference coordinate (second site)
253          * \param t simulation time
254          * \return calculated force and energy
255          *
256          * \todo remove to template header.
257          */
258         gmx::PotentialPointData evaluate(gmx::Vector r1,
259                                          gmx::Vector r2,
260                                          double t) override
261         {
262             return calculate(r1,
263                              r2,
264                              t);
265         };
266
267         /*!
268          * \brief An update function to be called on the simulation master rank/thread periodically by the Restraint framework.
269          *
270          * Implements optional override of gmx::IRestraintPotential::update
271          *
272          * This boilerplate will disappear into the Restraint template in an upcoming gmxapi release.
273          */
274         void update(gmx::Vector v,
275                     gmx::Vector v0,
276                     double t) override
277         {
278             // Todo: use a callback period to mostly bypass this and avoid excessive mutex locking.
279             callback(v,
280                      v0,
281                      t,
282                      *resources_);
283         };
284
285         /*!
286          * \brief Implement the binding protocol that allows access to Session resources.
287          *
288          * The client receives a non-owning pointer to the session and cannot extent the life of the session. In
289          * the future we can use a more formal handle mechanism.
290          *
291          * \param session pointer to the current session
292          */
293         void bindSession(gmxapi::SessionResources* session) override
294         {
295             resources_->setSession(session);
296         }
297
298         void setResources(std::unique_ptr<Resources>&& resources)
299         {
300             resources_ = std::move(resources);
301         }
302
303     private:
304         std::vector<int> sites_;
305 //        double callbackPeriod_;
306 //        double nextCallback_;
307         std::shared_ptr<Resources> resources_;
308 };
309
310
311 // Important: Just declare the template instantiation here for client code.
312 // We will explicitly instantiate a definition in the .cpp file where the input_param_type is defined.
313 extern template
314 class RestraintModule<EnsembleRestraint>;
315
316 } // end namespace plugin
317
318 #endif //HARMONICRESTRAINT_ENSEMBLEPOTENTIAL_H