0ea52a0b8d054a779253b09c46c96b608c29a7f8
[alexxy/gromacs.git] / src / gromacs / analysisdata / modules / plot.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2010-2018, The GROMACS development team.
5  * Copyright (c) 2019,2020, by the GROMACS development team, led by
6  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
7  * and including many others, as listed in the AUTHORS file in the
8  * top-level source directory and at http://www.gromacs.org.
9  *
10  * GROMACS is free software; you can redistribute it and/or
11  * modify it under the terms of the GNU Lesser General Public License
12  * as published by the Free Software Foundation; either version 2.1
13  * of the License, or (at your option) any later version.
14  *
15  * GROMACS is distributed in the hope that it will be useful,
16  * but WITHOUT ANY WARRANTY; without even the implied warranty of
17  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
18  * Lesser General Public License for more details.
19  *
20  * You should have received a copy of the GNU Lesser General Public
21  * License along with GROMACS; if not, see
22  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
23  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
24  *
25  * If you want to redistribute modifications to GROMACS, please
26  * consider that scientific software is very special. Version
27  * control is crucial - bugs must be traceable. We will be happy to
28  * consider code for inclusion in the official distribution, but
29  * derived work must not be called official GROMACS. Details are found
30  * in the README & COPYING files - if they are missing, get the
31  * official version at http://www.gromacs.org.
32  *
33  * To help us fund GROMACS development, we humbly ask that you cite
34  * the research papers on the package. Check out http://www.gromacs.org.
35  */
36 /*! \internal \file
37  * \brief
38  * Implements classes in plot.h.
39  *
40  * \ingroup module_analysisdata
41  * \author Teemu Murtola <teemu.murtola@gmail.com>
42  */
43 #include "gmxpre.h"
44
45 #include "plot.h"
46
47 #include <cstdio>
48 #include <cstring>
49
50 #include <string>
51 #include <vector>
52
53 #include "gromacs/analysisdata/dataframe.h"
54 #include "gromacs/fileio/gmxfio.h"
55 #include "gromacs/fileio/oenv.h"
56 #include "gromacs/fileio/xvgr.h"
57 #include "gromacs/math/vec.h"
58 #include "gromacs/options/basicoptions.h"
59 #include "gromacs/options/ioptionscontainer.h"
60 #include "gromacs/options/timeunitmanager.h"
61 #include "gromacs/selection/selectioncollection.h"
62 #include "gromacs/utility/exceptions.h"
63 #include "gromacs/utility/gmxassert.h"
64 #include "gromacs/utility/programcontext.h"
65 #include "gromacs/utility/stringutil.h"
66 #include "gromacs/utility/unique_cptr.h"
67
68 namespace gmx
69 {
70
71 /********************************************************************
72  * AnalysisDataPlotSettings
73  */
74
75 AnalysisDataPlotSettings::AnalysisDataPlotSettings() :
76     selections_(nullptr),
77     timeUnit_(TimeUnit::Default),
78     plotFormat_(XvgFormat::Xmgrace)
79 {
80 }
81
82 void AnalysisDataPlotSettings::setSelectionCollection(const SelectionCollection* selections)
83 {
84     selections_ = selections;
85 }
86
87 /*! \brief Names for XvgFormat
88  *
89  * Technically this duplicates a definition in pargs.cpp for legacy
90  * support code, but as the latter will go away and the alternatives
91  * are ugly, the duplication is acceptable. */
92 const gmx::EnumerationArray<XvgFormat, const char*> c_xvgFormatNames = { { "xmgrace", "xmgr",
93                                                                            "none" } };
94
95 void AnalysisDataPlotSettings::initOptions(IOptionsContainer* options)
96 {
97     options->addOption(
98             EnumOption<XvgFormat>("xvg").enumValue(c_xvgFormatNames).store(&plotFormat_).description("Plot formatting"));
99 }
100
101
102 /********************************************************************
103  * AbstractPlotModule::Impl
104  */
105
106 class AbstractPlotModule::Impl
107 {
108 public:
109     explicit Impl(const AnalysisDataPlotSettings& settings);
110     ~Impl();
111
112     void closeFile();
113
114     AnalysisDataPlotSettings settings_;
115     std::string              filename_;
116     FILE*                    fp_;
117
118     bool                     bPlain_;
119     bool                     bOmitX_;
120     bool                     bErrorsAsSeparateColumn_;
121     std::string              title_;
122     std::string              subtitle_;
123     std::string              xlabel_;
124     std::string              ylabel_;
125     std::vector<std::string> legend_;
126     std::string              xformat_;
127     std::string              yformat_;
128     real                     xscale_;
129 };
130
131 AbstractPlotModule::Impl::Impl(const AnalysisDataPlotSettings& settings) :
132     settings_(settings),
133     fp_(nullptr),
134     bPlain_(false),
135     bOmitX_(false),
136     bErrorsAsSeparateColumn_(false),
137     xformat_("%11.3f"),
138     yformat_(" %8.3f"),
139     xscale_(1.0)
140 {
141 }
142
143 AbstractPlotModule::Impl::~Impl()
144 {
145     closeFile();
146 }
147
148
149 void AbstractPlotModule::Impl::closeFile()
150 {
151     if (fp_ != nullptr)
152     {
153         if (bPlain_)
154         {
155             gmx_fio_fclose(fp_);
156         }
157         else
158         {
159             xvgrclose(fp_);
160         }
161         fp_ = nullptr;
162     }
163 }
164
165
166 /********************************************************************
167  * AbstractPlotModule
168  */
169 /*! \cond libapi */
170 AbstractPlotModule::AbstractPlotModule() : impl_(new Impl(AnalysisDataPlotSettings())) {}
171
172 AbstractPlotModule::AbstractPlotModule(const AnalysisDataPlotSettings& settings) :
173     impl_(new Impl(settings))
174 {
175 }
176 //! \endcond
177
178 AbstractPlotModule::~AbstractPlotModule() {}
179
180
181 void AbstractPlotModule::setSettings(const AnalysisDataPlotSettings& settings)
182 {
183     impl_->settings_ = settings;
184 }
185
186
187 void AbstractPlotModule::setFileName(const std::string& filename)
188 {
189     impl_->filename_ = filename;
190 }
191
192
193 void AbstractPlotModule::setPlainOutput(bool bPlain)
194 {
195     impl_->bPlain_ = bPlain;
196 }
197
198
199 void AbstractPlotModule::setErrorsAsSeparateColumn(bool bSeparate)
200 {
201     impl_->bErrorsAsSeparateColumn_ = bSeparate;
202 }
203
204
205 void AbstractPlotModule::setOmitX(bool bOmitX)
206 {
207     impl_->bOmitX_ = bOmitX;
208 }
209
210
211 void AbstractPlotModule::setTitle(const char* title)
212 {
213     impl_->title_ = title;
214 }
215
216 void AbstractPlotModule::setTitle(const std::string& title)
217 {
218     impl_->title_ = title;
219 }
220
221
222 void AbstractPlotModule::setSubtitle(const char* subtitle)
223 {
224     impl_->subtitle_ = subtitle;
225 }
226
227
228 void AbstractPlotModule::setSubtitle(const std::string& subtitle)
229 {
230     impl_->subtitle_ = subtitle;
231 }
232
233
234 void AbstractPlotModule::setXLabel(const char* label)
235 {
236     impl_->xlabel_ = label;
237 }
238
239
240 void AbstractPlotModule::setXAxisIsTime()
241 {
242     TimeUnitManager manager(impl_->settings_.timeUnit());
243     impl_->xlabel_ = formatString("Time (%s)", manager.timeUnitAsString());
244     impl_->xscale_ = manager.inverseTimeScaleFactor();
245 }
246
247
248 void AbstractPlotModule::setYLabel(const char* label)
249 {
250     impl_->ylabel_ = label;
251 }
252
253
254 void AbstractPlotModule::setLegend(int nsets, const char* const* setname)
255 {
256     impl_->legend_.reserve(impl_->legend_.size() + nsets);
257     for (int i = 0; i < nsets; ++i)
258     {
259         appendLegend(setname[i]);
260     }
261 }
262
263
264 void AbstractPlotModule::appendLegend(const char* setname)
265 {
266     impl_->legend_.emplace_back(setname);
267 }
268
269
270 void AbstractPlotModule::appendLegend(const std::string& setname)
271 {
272     impl_->legend_.push_back(setname);
273 }
274
275
276 void AbstractPlotModule::setXFormat(int width, int precision, char format)
277 {
278     GMX_RELEASE_ASSERT(width >= 0 && precision >= 0 && width <= 99 && precision <= 99,
279                        "Invalid width or precision");
280     GMX_RELEASE_ASSERT(strchr("eEfFgG", format) != nullptr, "Invalid format specifier");
281     impl_->xformat_ = formatString("%%%d.%d%c", width, precision, format);
282 }
283
284
285 void AbstractPlotModule::setYFormat(int width, int precision, char format)
286 {
287     GMX_RELEASE_ASSERT(width >= 0 && precision >= 0 && width <= 99 && precision <= 99,
288                        "Invalid width or precision");
289     GMX_RELEASE_ASSERT(strchr("eEfFgG", format) != nullptr, "Invalid format specifier");
290     impl_->yformat_ = formatString(" %%%d.%d%c", width, precision, format);
291 }
292
293
294 int AbstractPlotModule::flags() const
295 {
296     return efAllowMissing | efAllowMulticolumn | efAllowMultipoint | efAllowMultipleDataSets;
297 }
298
299
300 void AbstractPlotModule::dataStarted(AbstractAnalysisData* /* data */)
301 {
302     if (!impl_->filename_.empty())
303     {
304         if (impl_->bPlain_)
305         {
306             impl_->fp_ = gmx_fio_fopen(impl_->filename_.c_str(), "w");
307         }
308         else
309         {
310             const TimeUnit    timeUnit  = impl_->settings_.timeUnit();
311             const XvgFormat   xvgFormat = impl_->settings_.plotFormat();
312             gmx_output_env_t* oenv;
313             output_env_init(&oenv, getProgramContext(), timeUnit, FALSE, xvgFormat, 0);
314             const unique_cptr<gmx_output_env_t, output_env_done> oenvGuard(oenv);
315             impl_->fp_ = xvgropen(impl_->filename_.c_str(), impl_->title_.c_str(), impl_->xlabel_,
316                                   impl_->ylabel_, oenv);
317             const SelectionCollection* selections = impl_->settings_.selectionCollection();
318             if (selections != nullptr && output_env_get_xvg_format(oenv) != XvgFormat::None)
319             {
320                 selections->printXvgrInfo(impl_->fp_);
321             }
322             if (!impl_->subtitle_.empty())
323             {
324                 xvgr_subtitle(impl_->fp_, impl_->subtitle_.c_str(), oenv);
325             }
326             if (output_env_get_print_xvgr_codes(oenv) && !impl_->legend_.empty())
327             {
328                 std::vector<const char*> legend;
329                 legend.reserve(impl_->legend_.size());
330                 for (size_t i = 0; i < impl_->legend_.size(); ++i)
331                 {
332                     legend.push_back(impl_->legend_[i].c_str());
333                 }
334                 xvgr_legend(impl_->fp_, legend.size(), legend.data(), oenv);
335             }
336         }
337     }
338 }
339
340
341 void AbstractPlotModule::frameStarted(const AnalysisDataFrameHeader& header)
342 {
343     if (!isFileOpen())
344     {
345         return;
346     }
347     if (!impl_->bOmitX_)
348     {
349         std::fprintf(impl_->fp_, impl_->xformat_.c_str(), header.x() * impl_->xscale_);
350     }
351 }
352
353
354 void AbstractPlotModule::frameFinished(const AnalysisDataFrameHeader& /*header*/)
355 {
356     if (!isFileOpen())
357     {
358         return;
359     }
360     std::fprintf(impl_->fp_, "\n");
361 }
362
363
364 void AbstractPlotModule::dataFinished()
365 {
366     impl_->closeFile();
367 }
368
369 /*! \cond libapi */
370 bool AbstractPlotModule::isFileOpen() const
371 {
372     return impl_->fp_ != nullptr;
373 }
374
375
376 void AbstractPlotModule::writeValue(const AnalysisDataValue& value) const
377 {
378     GMX_ASSERT(isFileOpen(), "File not opened, but write attempted");
379     const real y = value.isSet() ? value.value() : 0.0;
380     std::fprintf(impl_->fp_, impl_->yformat_.c_str(), y);
381     if (impl_->bErrorsAsSeparateColumn_)
382     {
383         const real dy = value.isSet() ? value.error() : 0.0;
384         std::fprintf(impl_->fp_, impl_->yformat_.c_str(), dy);
385     }
386 }
387 //! \endcond
388
389 /********************************************************************
390  * DataPlotModule
391  */
392
393 AnalysisDataPlotModule::AnalysisDataPlotModule() {}
394
395 AnalysisDataPlotModule::AnalysisDataPlotModule(const AnalysisDataPlotSettings& settings) :
396     AbstractPlotModule(settings)
397 {
398 }
399
400
401 void AnalysisDataPlotModule::pointsAdded(const AnalysisDataPointSetRef& points)
402 {
403     if (!isFileOpen())
404     {
405         return;
406     }
407     for (int i = 0; i < points.columnCount(); ++i)
408     {
409         writeValue(points.values()[i]);
410     }
411 }
412
413
414 /********************************************************************
415  * DataVectorPlotModule
416  */
417
418 AnalysisDataVectorPlotModule::AnalysisDataVectorPlotModule() : bWrite_{ true, true, true, false } {}
419
420
421 AnalysisDataVectorPlotModule::AnalysisDataVectorPlotModule(const AnalysisDataPlotSettings& settings) :
422     AbstractPlotModule(settings),
423     bWrite_{ true, true, true, false }
424 {
425 }
426
427
428 void AnalysisDataVectorPlotModule::setWriteX(bool bWrite)
429 {
430     bWrite_[XX] = bWrite;
431 }
432
433
434 void AnalysisDataVectorPlotModule::setWriteY(bool bWrite)
435 {
436     bWrite_[YY] = bWrite;
437 }
438
439
440 void AnalysisDataVectorPlotModule::setWriteZ(bool bWrite)
441 {
442     bWrite_[ZZ] = bWrite;
443 }
444
445
446 void AnalysisDataVectorPlotModule::setWriteNorm(bool bWrite)
447 {
448     bWrite_[DIM] = bWrite;
449 }
450
451
452 void AnalysisDataVectorPlotModule::setWriteMask(const bool bWrite[DIM + 1])
453 {
454     for (int i = 0; i < DIM + 1; ++i)
455     {
456         bWrite_[i] = bWrite[i];
457     }
458 }
459
460
461 void AnalysisDataVectorPlotModule::pointsAdded(const AnalysisDataPointSetRef& points)
462 {
463     if (points.firstColumn() % DIM != 0 || points.columnCount() % DIM != 0)
464     {
465         GMX_THROW(APIError("Partial data points"));
466     }
467     if (!isFileOpen())
468     {
469         return;
470     }
471     for (int i = 0; i < points.columnCount(); i += 3)
472     {
473         for (int d = 0; d < DIM; ++d)
474         {
475             if (bWrite_[d])
476             {
477                 writeValue(points.values()[i + d]);
478             }
479         }
480         if (bWrite_[DIM])
481         {
482             const rvec        y = { points.y(i), points.y(i + 1), points.y(i + 2) };
483             AnalysisDataValue value(norm(y));
484             writeValue(value);
485         }
486     }
487 }
488
489 } // namespace gmx