c80a5bba386379618be33126efa255967680ebe8
[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,2011,2012,2013,2014,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 /*! \internal \file
36  * \brief
37  * Implements classes in plot.h.
38  *
39  * \ingroup module_analysisdata
40  * \author Teemu Murtola <teemu.murtola@gmail.com>
41  */
42 #include "gmxpre.h"
43
44 #include "plot.h"
45
46 #include <cstdio>
47 #include <cstring>
48
49 #include <string>
50 #include <vector>
51
52 #include "gromacs/analysisdata/dataframe.h"
53 #include "gromacs/fileio/gmxfio.h"
54 #include "gromacs/fileio/oenv.h"
55 #include "gromacs/fileio/xvgr.h"
56 #include "gromacs/math/vec.h"
57 #include "gromacs/options/basicoptions.h"
58 #include "gromacs/options/ioptionscontainer.h"
59 #include "gromacs/options/timeunitmanager.h"
60 #include "gromacs/selection/selectioncollection.h"
61 #include "gromacs/utility/exceptions.h"
62 #include "gromacs/utility/gmxassert.h"
63 #include "gromacs/utility/programcontext.h"
64 #include "gromacs/utility/stringutil.h"
65 #include "gromacs/utility/unique_cptr.h"
66
67 namespace
68 {
69
70 //! Enum values for plot formats.
71 const char *const g_plotFormats[] = {
72     "none", "xmgrace", "xmgr"
73 };
74
75 } // namespace
76
77 namespace gmx
78 {
79
80 /********************************************************************
81  * AnalysisDataPlotSettings
82  */
83
84 AnalysisDataPlotSettings::AnalysisDataPlotSettings()
85     : selections_(nullptr), timeUnit_(TimeUnit_Default), plotFormat_(1)
86 {
87 }
88
89 void
90 AnalysisDataPlotSettings::setSelectionCollection(const SelectionCollection *selections)
91 {
92     selections_ = selections;
93 }
94
95
96 void
97 AnalysisDataPlotSettings::initOptions(IOptionsContainer *options)
98 {
99     options->addOption(EnumIntOption("xvg").enumValue(g_plotFormats)
100                            .store(&plotFormat_)
101                            .description("Plot formatting"));
102 }
103
104
105 /********************************************************************
106  * AbstractPlotModule::Impl
107  */
108
109 class AbstractPlotModule::Impl
110 {
111     public:
112         explicit Impl(const AnalysisDataPlotSettings &settings);
113         ~Impl();
114
115         void closeFile();
116
117         AnalysisDataPlotSettings  settings_;
118         std::string               filename_;
119         FILE                     *fp_;
120
121         bool                      bPlain_;
122         bool                      bOmitX_;
123         bool                      bErrorsAsSeparateColumn_;
124         std::string               title_;
125         std::string               subtitle_;
126         std::string               xlabel_;
127         std::string               ylabel_;
128         std::vector<std::string>  legend_;
129         char                      xformat_[15];
130         char                      yformat_[15];
131         real                      xscale_;
132 };
133
134 AbstractPlotModule::Impl::Impl(const AnalysisDataPlotSettings &settings)
135     : settings_(settings), fp_(nullptr), bPlain_(false), bOmitX_(false),
136       bErrorsAsSeparateColumn_(false), xscale_(1.0)
137 {
138     strcpy(xformat_, "%11.3f");
139     strcpy(yformat_, " %8.3f");
140 }
141
142 AbstractPlotModule::Impl::~Impl()
143 {
144     closeFile();
145 }
146
147
148 void
149 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()
171     : impl_(new Impl(AnalysisDataPlotSettings()))
172 {
173 }
174
175 AbstractPlotModule::AbstractPlotModule(const AnalysisDataPlotSettings &settings)
176     : impl_(new Impl(settings))
177 {
178 }
179 //! \endcond
180
181 AbstractPlotModule::~AbstractPlotModule()
182 {
183 }
184
185
186 void
187 AbstractPlotModule::setSettings(const AnalysisDataPlotSettings &settings)
188 {
189     impl_->settings_ = settings;
190 }
191
192
193 void
194 AbstractPlotModule::setFileName(const std::string &filename)
195 {
196     impl_->filename_ = filename;
197 }
198
199
200 void
201 AbstractPlotModule::setPlainOutput(bool bPlain)
202 {
203     impl_->bPlain_ = bPlain;
204 }
205
206
207 void
208 AbstractPlotModule::setErrorsAsSeparateColumn(bool bSeparate)
209 {
210     impl_->bErrorsAsSeparateColumn_ = bSeparate;
211 }
212
213
214 void
215 AbstractPlotModule::setOmitX(bool bOmitX)
216 {
217     impl_->bOmitX_ = bOmitX;
218 }
219
220
221 void
222 AbstractPlotModule::setTitle(const char *title)
223 {
224     impl_->title_ = title;
225 }
226
227 void
228 AbstractPlotModule::setTitle(const std::string &title)
229 {
230     impl_->title_ = title;
231 }
232
233
234 void
235 AbstractPlotModule::setSubtitle(const char *subtitle)
236 {
237     impl_->subtitle_ = subtitle;
238 }
239
240
241 void
242 AbstractPlotModule::setSubtitle(const std::string &subtitle)
243 {
244     impl_->subtitle_ = subtitle;
245 }
246
247
248 void
249 AbstractPlotModule::setXLabel(const char *label)
250 {
251     impl_->xlabel_ = label;
252 }
253
254
255 void
256 AbstractPlotModule::setXAxisIsTime()
257 {
258     TimeUnitManager manager(impl_->settings_.timeUnit());
259     impl_->xlabel_ = formatString("Time (%s)", manager.timeUnitAsString());
260     impl_->xscale_ = manager.inverseTimeScaleFactor();
261 }
262
263
264 void
265 AbstractPlotModule::setYLabel(const char *label)
266 {
267     impl_->ylabel_ = label;
268 }
269
270
271 void
272 AbstractPlotModule::setLegend(int nsets, const char * const *setname)
273 {
274     impl_->legend_.reserve(impl_->legend_.size() + nsets);
275     for (int i = 0; i < nsets; ++i)
276     {
277         appendLegend(setname[i]);
278     }
279 }
280
281
282 void
283 AbstractPlotModule::appendLegend(const char *setname)
284 {
285     impl_->legend_.emplace_back(setname);
286 }
287
288
289 void
290 AbstractPlotModule::appendLegend(const std::string &setname)
291 {
292     impl_->legend_.push_back(setname);
293 }
294
295
296 void
297 AbstractPlotModule::setXFormat(int width, int precision, char format)
298 {
299     GMX_RELEASE_ASSERT(width >= 0 && precision >= 0
300                        && width <= 99 && precision <= 99,
301                        "Invalid width or precision");
302     GMX_RELEASE_ASSERT(strchr("eEfFgG", format) != nullptr,
303                        "Invalid format specifier");
304     sprintf(impl_->xformat_, "%%%d.%d%c", width, precision, format);
305 }
306
307
308 void
309 AbstractPlotModule::setYFormat(int width, int precision, char format)
310 {
311     GMX_RELEASE_ASSERT(width >= 0 && precision >= 0
312                        && width <= 99 && precision <= 99,
313                        "Invalid width or precision");
314     GMX_RELEASE_ASSERT(strchr("eEfFgG", format) != nullptr,
315                        "Invalid format specifier");
316     sprintf(impl_->yformat_, " %%%d.%d%c", width, precision, format);
317 }
318
319
320 int
321 AbstractPlotModule::flags() const
322 {
323     return efAllowMissing | efAllowMulticolumn | efAllowMultipoint
324            | efAllowMultipleDataSets;
325 }
326
327
328 void
329 AbstractPlotModule::dataStarted(AbstractAnalysisData * /* data */)
330 {
331     if (!impl_->filename_.empty())
332     {
333         if (impl_->bPlain_)
334         {
335             impl_->fp_ = gmx_fio_fopen(impl_->filename_.c_str(), "w");
336         }
337         else
338         {
339             time_unit_t  time_unit
340                 = static_cast<time_unit_t>(impl_->settings_.timeUnit() + 1);
341             xvg_format_t xvg_format
342                 = (impl_->settings_.plotFormat() > 0
343                    ? static_cast<xvg_format_t>(impl_->settings_.plotFormat())
344                    : exvgNONE);
345             gmx_output_env_t                                    *oenv;
346             output_env_init(&oenv, getProgramContext(), time_unit, FALSE, xvg_format, 0);
347             const unique_cptr<gmx_output_env_t, output_env_done> oenvGuard(oenv);
348             impl_->fp_ = xvgropen(impl_->filename_.c_str(), impl_->title_.c_str(),
349                                   impl_->xlabel_, impl_->ylabel_,
350                                   oenv);
351             const SelectionCollection *selections
352                 = impl_->settings_.selectionCollection();
353             if (selections != nullptr && output_env_get_xvg_format(oenv) != exvgNONE)
354             {
355                 selections->printXvgrInfo(impl_->fp_);
356             }
357             if (!impl_->subtitle_.empty())
358             {
359                 xvgr_subtitle(impl_->fp_, impl_->subtitle_.c_str(), oenv);
360             }
361             if (output_env_get_print_xvgr_codes(oenv)
362                 && !impl_->legend_.empty())
363             {
364                 std::vector<const char *> legend;
365                 legend.reserve(impl_->legend_.size());
366                 for (size_t i = 0; i < impl_->legend_.size(); ++i)
367                 {
368                     legend.push_back(impl_->legend_[i].c_str());
369                 }
370                 xvgr_legend(impl_->fp_, legend.size(), legend.data(), oenv);
371             }
372         }
373     }
374 }
375
376
377 void
378 AbstractPlotModule::frameStarted(const AnalysisDataFrameHeader &frame)
379 {
380     if (!isFileOpen())
381     {
382         return;
383     }
384     if (!impl_->bOmitX_)
385     {
386         std::fprintf(impl_->fp_, impl_->xformat_, frame.x() * impl_->xscale_);
387     }
388 }
389
390
391 void
392 AbstractPlotModule::frameFinished(const AnalysisDataFrameHeader & /*header*/)
393 {
394     if (!isFileOpen())
395     {
396         return;
397     }
398     std::fprintf(impl_->fp_, "\n");
399 }
400
401
402 void
403 AbstractPlotModule::dataFinished()
404 {
405     impl_->closeFile();
406 }
407
408 /*! \cond libapi */
409 bool
410 AbstractPlotModule::isFileOpen() const
411 {
412     return impl_->fp_ != nullptr;
413 }
414
415
416 void
417 AbstractPlotModule::writeValue(const AnalysisDataValue &value) const
418 {
419     GMX_ASSERT(isFileOpen(), "File not opened, but write attempted");
420     const real y = value.isSet() ? value.value() : 0.0;
421     std::fprintf(impl_->fp_, impl_->yformat_, y);
422     if (impl_->bErrorsAsSeparateColumn_)
423     {
424         const real dy = value.isSet() ? value.error() : 0.0;
425         std::fprintf(impl_->fp_, impl_->yformat_, dy);
426     }
427 }
428 //! \endcond
429
430 /********************************************************************
431  * DataPlotModule
432  */
433
434 AnalysisDataPlotModule::AnalysisDataPlotModule()
435 {
436 }
437
438 AnalysisDataPlotModule::AnalysisDataPlotModule(
439         const AnalysisDataPlotSettings &settings)
440     : AbstractPlotModule(settings)
441 {
442 }
443
444
445 void
446 AnalysisDataPlotModule::pointsAdded(const AnalysisDataPointSetRef &points)
447 {
448     if (!isFileOpen())
449     {
450         return;
451     }
452     for (int i = 0; i < points.columnCount(); ++i)
453     {
454         writeValue(points.values()[i]);
455     }
456 }
457
458
459 /********************************************************************
460  * DataVectorPlotModule
461  */
462
463 AnalysisDataVectorPlotModule::AnalysisDataVectorPlotModule()
464 {
465     for (int i = 0; i < DIM; ++i)
466     {
467         bWrite_[i] = true;
468     }
469     bWrite_[DIM] = false;
470 }
471
472
473 AnalysisDataVectorPlotModule::AnalysisDataVectorPlotModule(
474         const AnalysisDataPlotSettings &settings)
475     : AbstractPlotModule(settings)
476 {
477     for (int i = 0; i < DIM; ++i)
478     {
479         bWrite_[i] = true;
480     }
481     bWrite_[DIM] = false;
482 }
483
484
485 void
486 AnalysisDataVectorPlotModule::setWriteX(bool bWrite)
487 {
488     bWrite_[XX] = bWrite;
489 }
490
491
492 void
493 AnalysisDataVectorPlotModule::setWriteY(bool bWrite)
494 {
495     bWrite_[YY] = bWrite;
496 }
497
498
499 void
500 AnalysisDataVectorPlotModule::setWriteZ(bool bWrite)
501 {
502     bWrite_[ZZ] = bWrite;
503 }
504
505
506 void
507 AnalysisDataVectorPlotModule::setWriteNorm(bool bWrite)
508 {
509     bWrite_[DIM] = bWrite;
510 }
511
512
513 void
514 AnalysisDataVectorPlotModule::setWriteMask(bool bWrite[DIM + 1])
515 {
516     for (int i = 0; i < DIM + 1; ++i)
517     {
518         bWrite_[i] = bWrite[i];
519     }
520 }
521
522
523 void
524 AnalysisDataVectorPlotModule::pointsAdded(const AnalysisDataPointSetRef &points)
525 {
526     if (points.firstColumn() % DIM != 0 || points.columnCount() % DIM != 0)
527     {
528         GMX_THROW(APIError("Partial data points"));
529     }
530     if (!isFileOpen())
531     {
532         return;
533     }
534     for (int i = 0; i < points.columnCount(); i += 3)
535     {
536         for (int d = 0; d < DIM; ++d)
537         {
538             if (bWrite_[d])
539             {
540                 writeValue(points.values()[i + d]);
541             }
542         }
543         if (bWrite_[DIM])
544         {
545             const rvec        y = { points.y(i), points.y(i + 1), points.y(i + 2) };
546             AnalysisDataValue value(norm(y));
547             writeValue(value);
548         }
549     }
550 }
551
552 } // namespace gmx