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