38b8bb67ea73cf0132c5a972db564988fb8662c7
[alexxy/gromacs.git] / src / gromacs / selection / poscalc.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2009,2010,2011,2012,2013 by the GROMACS development team.
5  * Copyright (c) 2014,2015,2016,2017,2018 by the GROMACS development team.
6  * Copyright (c) 2019,2020, by the GROMACS development team, led by
7  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
8  * and including many others, as listed in the AUTHORS file in the
9  * top-level source directory and at http://www.gromacs.org.
10  *
11  * GROMACS is free software; you can redistribute it and/or
12  * modify it under the terms of the GNU Lesser General Public License
13  * as published by the Free Software Foundation; either version 2.1
14  * of the License, or (at your option) any later version.
15  *
16  * GROMACS is distributed in the hope that it will be useful,
17  * but WITHOUT ANY WARRANTY; without even the implied warranty of
18  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
19  * Lesser General Public License for more details.
20  *
21  * You should have received a copy of the GNU Lesser General Public
22  * License along with GROMACS; if not, see
23  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
24  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
25  *
26  * If you want to redistribute modifications to GROMACS, please
27  * consider that scientific software is very special. Version
28  * control is crucial - bugs must be traceable. We will be happy to
29  * consider code for inclusion in the official distribution, but
30  * derived work must not be called official GROMACS. Details are found
31  * in the README & COPYING files - if they are missing, get the
32  * official version at http://www.gromacs.org.
33  *
34  * To help us fund GROMACS development, we humbly ask that you cite
35  * the research papers on the package. Check out http://www.gromacs.org.
36  */
37 /*! \internal \file
38  * \brief
39  * Implements gmx::PositionCalculationCollection and functions in poscalc.h.
40  *
41  * \todo
42  * There is probably some room for optimization in the calculation of
43  * positions with bases.
44  * In particular, the current implementation may do a lot of unnecessary
45  * copying.
46  * The interface would need to be changed to make it possible to use the
47  * same output positions for several calculations.
48  *
49  * \todo
50  * The current algorithm for setting up base calculations could be improved
51  * in cases when there are calculations that cannot use a common base but
52  * still overlap partially (e.g., with three calculations A, B, and C
53  * such that A could use both B and C as a base, but B and C cannot use the
54  * same base).
55  * Setting up the bases in an optimal manner in every possible situation can
56  * be quite difficult unless several bases are allowed for one calculation,
57  * but better heuristics could probably be implemented.
58  * For best results, the setup should probably be postponed (at least
59  * partially) to gmx_ana_poscalc_init_eval().
60  *
61  * \author Teemu Murtola <teemu.murtola@gmail.com>
62  * \ingroup module_selection
63  */
64 #include "gmxpre.h"
65
66 #include "poscalc.h"
67
68 #include <cstring>
69
70 #include <algorithm>
71 #include <vector>
72
73 #include "gromacs/math/vec.h"
74 #include "gromacs/selection/indexutil.h"
75 #include "gromacs/trajectory/trajectoryframe.h"
76 #include "gromacs/utility/arrayref.h"
77 #include "gromacs/utility/exceptions.h"
78 #include "gromacs/utility/gmxassert.h"
79 #include "gromacs/utility/smalloc.h"
80
81 #include "centerofmass.h"
82 #include "position.h"
83
84 namespace gmx
85 {
86
87 /*! \internal \brief
88  * Private implementation class for PositionCalculationCollection.
89  *
90  * \ingroup module_selection
91  */
92 class PositionCalculationCollection::Impl
93 {
94 public:
95     Impl();
96     ~Impl();
97
98     /*! \brief
99      * Inserts a position calculation structure into its collection.
100      *
101      * \param pc     Data structure to insert.
102      * \param before Data structure before which to insert
103      *   (NULL = insert at end).
104      *
105      * Inserts \p pc to its collection before \p before.
106      * If \p before is NULL, \p pc is appended to the list.
107      */
108     void insertCalculation(gmx_ana_poscalc_t* pc, gmx_ana_poscalc_t* before);
109     /*! \brief
110      * Removes a position calculation structure from its collection.
111      *
112      * \param pc    Data structure to remove.
113      *
114      * Removes \p pc from its collection.
115      */
116     void removeCalculation(gmx_ana_poscalc_t* pc);
117
118     /*! \copydoc PositionCalculationCollection::createCalculation()
119      *
120      * This method contains the actual implementation of the similarly
121      * named method in the parent class.
122      */
123     gmx_ana_poscalc_t* createCalculation(e_poscalc_t type, int flags);
124
125     /*! \brief
126      * Maps given topology indices into frame indices.
127      *
128      * Only one position calculation at a time needs to access this (and
129      * there are also other thread-unsafe constructs here), so a temporary
130      * array is used to avoid repeated memory allocation.
131      */
132     ArrayRef<const int> getFrameIndices(int size, int index[])
133     {
134         if (mapToFrameAtoms_.empty())
135         {
136             return constArrayRefFromArray(index, size);
137         }
138         tmpFrameAtoms_.resize(size);
139         for (int i = 0; i < size; ++i)
140         {
141             const int ii = index[i];
142             GMX_ASSERT(ii >= 0 && ii <= ssize(mapToFrameAtoms_) && mapToFrameAtoms_[ii] != -1,
143                        "Invalid input atom index");
144             tmpFrameAtoms_[i] = mapToFrameAtoms_[ii];
145         }
146         return tmpFrameAtoms_;
147     }
148
149     /*! \brief
150      * Topology data.
151      *
152      * Can be NULL if none of the calculations require topology data or if
153      * setTopology() has not been called.
154      */
155     const gmx_mtop_t* top_;
156     //! Pointer to the first data structure.
157     gmx_ana_poscalc_t* first_;
158     //! Pointer to the last data structure.
159     gmx_ana_poscalc_t* last_;
160     //! Whether the collection has been initialized for evaluation.
161     bool bInit_;
162     //! Mapping from topology atoms to frame atoms (one index for each topology atom).
163     std::vector<int> mapToFrameAtoms_;
164     //! Working array for updating positions.
165     std::vector<int> tmpFrameAtoms_;
166 };
167
168 } // namespace gmx
169
170 /*! \internal \brief
171  * Data structure for position calculation.
172  */
173 struct gmx_ana_poscalc_t
174 {
175     /*! \brief
176      * Type of calculation.
177      *
178      * This field may differ from the type requested by the user, because
179      * it is changed internally to the most effective calculation.
180      * For example, if the user requests a COM calculation for residues
181      * consisting of single atoms, it is simply set to POS_ATOM.
182      * To provide a consistent interface to the user, the field \p itype
183      * should be used when information should be given out.
184      */
185     e_poscalc_t type;
186     /*! \brief
187      * Flags for calculation options.
188      *
189      * See \ref poscalc_flags "documentation of the flags".
190      */
191     int flags;
192
193     /*! \brief
194      * Type for the created indices.
195      *
196      * This field always agrees with the type that the user requested, but
197      * may differ from \p type.
198      */
199     e_index_t itype;
200     /*! \brief
201      * Block data for the calculation.
202      */
203     t_blocka b;
204     /*! \brief
205      * Mapping from the blocks to the blocks of \p sbase.
206      *
207      * If \p sbase is NULL, this field is also.
208      */
209     int* baseid;
210     /*! \brief
211      * Maximum evaluation group.
212      */
213     gmx_ana_index_t gmax;
214
215     /** Position storage for calculations that are used as a base. */
216     gmx_ana_pos_t* p;
217
218     /** true if the positions have been evaluated for the current frame. */
219     bool bEval;
220     /*! \brief
221      * Base position data for this calculation.
222      *
223      * If not NULL, the centers required by this calculation have
224      * already been calculated in \p sbase.
225      * The structure pointed by \p sbase is always a static calculation.
226      */
227     gmx_ana_poscalc_t* sbase;
228     /** Next structure in the linked list of calculations. */
229     gmx_ana_poscalc_t* next;
230     /** Previous structure in the linked list of calculations. */
231     gmx_ana_poscalc_t* prev;
232     /** Number of references to this structure. */
233     int refcount;
234     /** Collection this calculation belongs to. */
235     gmx::PositionCalculationCollection::Impl* coll;
236 };
237
238 const char* const gmx::PositionCalculationCollection::typeEnumValues[] = {
239     "atom",          "res_com",       "res_cog",       "mol_com",       "mol_cog",
240     "whole_res_com", "whole_res_cog", "whole_mol_com", "whole_mol_cog", "part_res_com",
241     "part_res_cog",  "part_mol_com",  "part_mol_cog",  "dyn_res_com",   "dyn_res_cog",
242     "dyn_mol_com",   "dyn_mol_cog",   nullptr,
243 };
244
245 /*! \brief
246  * Returns the partition type for a given position type.
247  *
248  * \param [in] type  \c e_poscalc_t value to convert.
249  * \returns    Corresponding \c e_indet_t.
250  */
251 static e_index_t index_type_for_poscalc(e_poscalc_t type)
252 {
253     switch (type)
254     {
255         case POS_ATOM: return INDEX_ATOM;
256         case POS_RES: return INDEX_RES;
257         case POS_MOL: return INDEX_MOL;
258         case POS_ALL: // Intended fall through
259         case POS_ALL_PBC: return INDEX_ALL;
260     }
261     return INDEX_UNKNOWN;
262 }
263
264 namespace gmx
265 {
266
267 namespace
268 {
269
270 //! Helper function for determining required topology information.
271 PositionCalculationCollection::RequiredTopologyInfo requiredTopologyInfo(e_poscalc_t type, int flags)
272 {
273     if (type != POS_ATOM)
274     {
275         if ((flags & POS_MASS) || (flags & POS_FORCES))
276         {
277             return PositionCalculationCollection::RequiredTopologyInfo::TopologyAndMasses;
278         }
279         if (type == POS_RES || type == POS_MOL)
280         {
281             return PositionCalculationCollection::RequiredTopologyInfo::Topology;
282         }
283     }
284     return PositionCalculationCollection::RequiredTopologyInfo::None;
285 }
286
287 } // namespace
288
289 // static
290 void PositionCalculationCollection::typeFromEnum(const char* post, e_poscalc_t* type, int* flags)
291 {
292     if (post[0] == 'a')
293     {
294         *type = POS_ATOM;
295         *flags &= ~(POS_MASS | POS_COMPLMAX | POS_COMPLWHOLE);
296         return;
297     }
298
299     /* Process the prefix */
300     const char* ptr = post;
301     if (post[0] == 'w')
302     {
303         *flags &= ~POS_COMPLMAX;
304         *flags |= POS_COMPLWHOLE;
305         ptr = post + 6;
306     }
307     else if (post[0] == 'p')
308     {
309         *flags &= ~POS_COMPLWHOLE;
310         *flags |= POS_COMPLMAX;
311         ptr = post + 5;
312     }
313     else if (post[0] == 'd')
314     {
315         *flags &= ~(POS_COMPLMAX | POS_COMPLWHOLE);
316         ptr = post + 4;
317     }
318
319     if (ptr[0] == 'r')
320     {
321         *type = POS_RES;
322     }
323     else if (ptr[0] == 'm')
324     {
325         *type = POS_MOL;
326     }
327     else
328     {
329         GMX_THROW(InternalError("Unknown position calculation type"));
330     }
331     if (strlen(ptr) < 7)
332     {
333         GMX_THROW(InternalError("Unknown position calculation type"));
334     }
335     if (ptr[6] == 'm')
336     {
337         *flags |= POS_MASS;
338     }
339     else if (ptr[6] == 'g')
340     {
341         *flags &= ~POS_MASS;
342     }
343     else
344     {
345         GMX_THROW(InternalError("Unknown position calculation type"));
346     }
347 }
348
349 // static
350 PositionCalculationCollection::RequiredTopologyInfo
351 PositionCalculationCollection::requiredTopologyInfoForType(const char* post, bool forces)
352 {
353     e_poscalc_t type;
354     int         flags = (forces ? POS_FORCES : 0);
355     PositionCalculationCollection::typeFromEnum(post, &type, &flags);
356     return requiredTopologyInfo(type, flags);
357 }
358
359 /********************************************************************
360  * PositionCalculationCollection::Impl
361  */
362
363 PositionCalculationCollection::Impl::Impl() :
364     top_(nullptr),
365     first_(nullptr),
366     last_(nullptr),
367     bInit_(false)
368 {
369 }
370
371 PositionCalculationCollection::Impl::~Impl()
372 {
373     // Loop backwards, because there can be internal references in that are
374     // correctly handled by this direction.
375     while (last_ != nullptr)
376     {
377         GMX_ASSERT(last_->refcount == 1, "Dangling references to position calculations");
378         gmx_ana_poscalc_free(last_);
379     }
380 }
381
382 void PositionCalculationCollection::Impl::insertCalculation(gmx_ana_poscalc_t* pc, gmx_ana_poscalc_t* before)
383 {
384     GMX_RELEASE_ASSERT(pc->coll == this, "Inconsistent collections");
385     if (before == nullptr)
386     {
387         pc->next = nullptr;
388         pc->prev = last_;
389         if (last_ != nullptr)
390         {
391             last_->next = pc;
392         }
393         last_ = pc;
394     }
395     else
396     {
397         pc->prev = before->prev;
398         pc->next = before;
399         if (before->prev)
400         {
401             before->prev->next = pc;
402         }
403         before->prev = pc;
404     }
405     if (pc->prev == nullptr)
406     {
407         first_ = pc;
408     }
409 }
410
411 void PositionCalculationCollection::Impl::removeCalculation(gmx_ana_poscalc_t* pc)
412 {
413     GMX_RELEASE_ASSERT(pc->coll == this, "Inconsistent collections");
414     if (pc->prev != nullptr)
415     {
416         pc->prev->next = pc->next;
417     }
418     else if (pc == first_)
419     {
420         first_ = pc->next;
421     }
422     if (pc->next != nullptr)
423     {
424         pc->next->prev = pc->prev;
425     }
426     else if (pc == last_)
427     {
428         last_ = pc->prev;
429     }
430     pc->prev = pc->next = nullptr;
431 }
432
433 gmx_ana_poscalc_t* PositionCalculationCollection::Impl::createCalculation(e_poscalc_t type, int flags)
434 {
435     gmx_ana_poscalc_t* pc;
436
437     snew(pc, 1);
438     pc->type  = type;
439     pc->itype = index_type_for_poscalc(type);
440     gmx_ana_poscalc_set_flags(pc, flags);
441     pc->refcount = 1;
442     pc->coll     = this;
443     insertCalculation(pc, nullptr);
444     return pc;
445 }
446
447
448 /********************************************************************
449  * PositionCalculationCollection
450  */
451
452 PositionCalculationCollection::PositionCalculationCollection() : impl_(new Impl) {}
453
454 PositionCalculationCollection::~PositionCalculationCollection() {}
455
456 void PositionCalculationCollection::setTopology(const gmx_mtop_t* top)
457 {
458     impl_->top_ = top;
459 }
460
461 void PositionCalculationCollection::printTree(FILE* fp) const
462 {
463     gmx_ana_poscalc_t* pc;
464     int                i, j;
465
466     fprintf(fp, "Position calculations:\n");
467     i  = 1;
468     pc = impl_->first_;
469     while (pc)
470     {
471         fprintf(fp, "%2d ", i);
472         switch (pc->type)
473         {
474             case POS_ATOM: fprintf(fp, "ATOM"); break;
475             case POS_RES: fprintf(fp, "RES"); break;
476             case POS_MOL: fprintf(fp, "MOL"); break;
477             case POS_ALL: fprintf(fp, "ALL"); break;
478             case POS_ALL_PBC: fprintf(fp, "ALL_PBC"); break;
479         }
480         if (pc->itype != index_type_for_poscalc(pc->type))
481         {
482             fprintf(fp, " (");
483             switch (pc->itype)
484             {
485                 case INDEX_UNKNOWN: fprintf(fp, "???"); break;
486                 case INDEX_ATOM: fprintf(fp, "ATOM"); break;
487                 case INDEX_RES: fprintf(fp, "RES"); break;
488                 case INDEX_MOL: fprintf(fp, "MOL"); break;
489                 case INDEX_ALL: fprintf(fp, "ALL"); break;
490             }
491             fprintf(fp, ")");
492         }
493         fprintf(fp, " flg=");
494         if (pc->flags & POS_MASS)
495         {
496             fprintf(fp, "M");
497         }
498         if (pc->flags & POS_DYNAMIC)
499         {
500             fprintf(fp, "D");
501         }
502         if (pc->flags & POS_MASKONLY)
503         {
504             fprintf(fp, "A");
505         }
506         if (pc->flags & POS_COMPLMAX)
507         {
508             fprintf(fp, "Cm");
509         }
510         if (pc->flags & POS_COMPLWHOLE)
511         {
512             fprintf(fp, "Cw");
513         }
514         if (!pc->flags)
515         {
516             fprintf(fp, "0");
517         }
518         fprintf(fp, " nr=%d nra=%d", pc->b.nr, pc->b.nra);
519         fprintf(fp, " refc=%d", pc->refcount);
520         fprintf(fp, "\n");
521         if (pc->gmax.nalloc_index > 0)
522         {
523             fprintf(fp, "   Group: ");
524             if (pc->gmax.isize > 20)
525             {
526                 fprintf(fp, " %d atoms", pc->gmax.isize);
527             }
528             else
529             {
530                 for (j = 0; j < pc->gmax.isize; ++j)
531                 {
532                     fprintf(fp, " %d", pc->gmax.index[j] + 1);
533                 }
534             }
535             fprintf(fp, "\n");
536         }
537         if (pc->b.nalloc_a > 0)
538         {
539             fprintf(fp, "   Atoms: ");
540             if (pc->b.nra > 20)
541             {
542                 fprintf(fp, " %d atoms", pc->b.nra);
543             }
544             else
545             {
546                 for (j = 0; j < pc->b.nra; ++j)
547                 {
548                     fprintf(fp, " %d", pc->b.a[j] + 1);
549                 }
550             }
551             fprintf(fp, "\n");
552         }
553         if (pc->b.nalloc_index > 0)
554         {
555             fprintf(fp, "   Blocks:");
556             if (pc->b.nr > 20)
557             {
558                 fprintf(fp, " %d pcs", pc->b.nr);
559             }
560             else
561             {
562                 for (j = 0; j <= pc->b.nr; ++j)
563                 {
564                     fprintf(fp, " %d", pc->b.index[j]);
565                 }
566             }
567             fprintf(fp, "\n");
568         }
569         if (pc->sbase)
570         {
571             gmx_ana_poscalc_t* base;
572
573             fprintf(fp, "   Base: ");
574             j    = 1;
575             base = impl_->first_;
576             while (base && base != pc->sbase)
577             {
578                 ++j;
579                 base = base->next;
580             }
581             fprintf(fp, "%d", j);
582             if (pc->baseid && pc->b.nr <= 20)
583             {
584                 fprintf(fp, " id:");
585                 for (j = 0; j < pc->b.nr; ++j)
586                 {
587                     fprintf(fp, " %d", pc->baseid[j] + 1);
588                 }
589             }
590             fprintf(fp, "\n");
591         }
592         ++i;
593         pc = pc->next;
594     }
595 }
596
597 gmx_ana_poscalc_t* PositionCalculationCollection::createCalculation(e_poscalc_t type, int flags)
598 {
599     return impl_->createCalculation(type, flags);
600 }
601
602 gmx_ana_poscalc_t* PositionCalculationCollection::createCalculationFromEnum(const char* post, int flags)
603 {
604     e_poscalc_t type;
605     int         cflags = flags;
606     typeFromEnum(post, &type, &cflags);
607     return impl_->createCalculation(type, cflags);
608 }
609
610 void PositionCalculationCollection::getRequiredAtoms(gmx_ana_index_t* out) const
611 {
612     gmx_ana_poscalc_t* pc = impl_->first_;
613     while (pc)
614     {
615         // Calculations with a base just copy positions from the base, so
616         // those do not need to be considered in the check.
617         if (!pc->sbase)
618         {
619             gmx_ana_index_t g;
620             gmx_ana_index_set(&g, pc->b.nra, pc->b.a, 0);
621             gmx_ana_index_union_unsorted(out, out, &g);
622         }
623         pc = pc->next;
624     }
625 }
626
627 void PositionCalculationCollection::initEvaluation()
628 {
629     if (impl_->bInit_)
630     {
631         return;
632     }
633     gmx_ana_poscalc_t* pc = impl_->first_;
634     while (pc)
635     {
636         /* Initialize position storage for base calculations */
637         if (pc->p)
638         {
639             gmx_ana_poscalc_init_pos(pc, pc->p);
640         }
641         /* Construct the mapping of the base positions */
642         if (pc->sbase)
643         {
644             int bi, bj;
645
646             snew(pc->baseid, pc->b.nr);
647             for (bi = bj = 0; bi < pc->b.nr; ++bi, ++bj)
648             {
649                 while (pc->sbase->b.a[pc->sbase->b.index[bj]] != pc->b.a[pc->b.index[bi]])
650                 {
651                     ++bj;
652                 }
653                 pc->baseid[bi] = bj;
654             }
655         }
656         /* Free the block data for dynamic calculations */
657         if (pc->flags & POS_DYNAMIC)
658         {
659             if (pc->b.nalloc_index > 0)
660             {
661                 sfree(pc->b.index);
662                 pc->b.nalloc_index = 0;
663             }
664             if (pc->b.nalloc_a > 0)
665             {
666                 sfree(pc->b.a);
667                 pc->b.nalloc_a = 0;
668             }
669         }
670         pc = pc->next;
671     }
672     impl_->bInit_ = true;
673 }
674
675 void PositionCalculationCollection::initFrame(const t_trxframe* fr)
676 {
677     if (!impl_->bInit_)
678     {
679         initEvaluation();
680     }
681     /* Clear the evaluation flags */
682     gmx_ana_poscalc_t* pc = impl_->first_;
683     while (pc)
684     {
685         pc->bEval = false;
686         pc        = pc->next;
687     }
688     if (fr->bIndex && fr->natoms > 0)
689     {
690         const int highestAtom = *std::max_element(fr->index, fr->index + fr->natoms);
691         impl_->mapToFrameAtoms_.resize(highestAtom + 1);
692         std::fill(impl_->mapToFrameAtoms_.begin(), impl_->mapToFrameAtoms_.end(), -1);
693         for (int i = 0; i < fr->natoms; ++i)
694         {
695             impl_->mapToFrameAtoms_[fr->index[i]] = i;
696         }
697     }
698     else
699     {
700         impl_->mapToFrameAtoms_.clear();
701     }
702 }
703
704 } // namespace gmx
705
706 /*! \brief
707  * Initializes position calculation using the maximum possible input index.
708  *
709  * \param[in,out] pc  Position calculation data structure.
710  * \param[in]     g   Maximum index group for the calculation.
711  * \param[in]     bBase Whether \p pc will be used as a base or not.
712  *
713  * \p bBase affects on how the \p pc->gmax field is initialized.
714  */
715 static void set_poscalc_maxindex(gmx_ana_poscalc_t* pc, gmx_ana_index_t* g, bool bBase)
716 {
717     const gmx_mtop_t* top = pc->coll->top_;
718     gmx_ana_index_make_block(&pc->b, top, g, pc->itype, (pc->flags & POS_COMPLWHOLE) != 0);
719     /* Set the type to POS_ATOM if the calculation in fact is such. */
720     if (pc->b.nr == pc->b.nra)
721     {
722         pc->type = POS_ATOM;
723         pc->flags &= ~(POS_MASS | POS_COMPLMAX | POS_COMPLWHOLE);
724     }
725     /* Set the POS_COMPLWHOLE flag if the calculation in fact always uses
726      * complete residues and molecules. */
727     if (!(pc->flags & POS_COMPLWHOLE) && (!(pc->flags & POS_DYNAMIC) || (pc->flags & POS_COMPLMAX))
728         && (pc->type == POS_RES || pc->type == POS_MOL)
729         && gmx_ana_index_has_complete_elems(g, pc->itype, top))
730     {
731         pc->flags &= ~POS_COMPLMAX;
732         pc->flags |= POS_COMPLWHOLE;
733     }
734     /* Setup the gmax field */
735     if ((pc->flags & POS_COMPLWHOLE) && !bBase && pc->b.nra > g->isize)
736     {
737         gmx_ana_index_copy(&pc->gmax, g, true);
738     }
739     else
740     {
741         gmx_ana_index_set(&pc->gmax, pc->b.nra, pc->b.a, 0);
742     }
743 }
744
745 /*! \brief
746  * Checks whether a position calculation should use a base at all.
747  *
748  * \param[in] pc   Position calculation data to check.
749  * \returns   true if \p pc can use a base and gets some benefit out of it,
750  *   false otherwise.
751  */
752 static bool can_use_base(gmx_ana_poscalc_t* pc)
753 {
754     /* For atoms, it should be faster to do a simple copy, so don't use a
755      * base. */
756     if (pc->type == POS_ATOM)
757     {
758         return false;
759     }
760     /* For dynamic selections that do not use completion, it is not possible
761      * to use a base. */
762     if ((pc->type == POS_RES || pc->type == POS_MOL) && (pc->flags & POS_DYNAMIC)
763         && !(pc->flags & (POS_COMPLMAX | POS_COMPLWHOLE)))
764     {
765         return false;
766     }
767     /* Dynamic calculations for a single position cannot use a base. */
768     if ((pc->type == POS_ALL || pc->type == POS_ALL_PBC) && (pc->flags & POS_DYNAMIC))
769     {
770         return false;
771     }
772     return true;
773 }
774
775 /*! \brief
776  * Checks whether two position calculations should use a common base.
777  *
778  * \param[in]     pc1 Calculation 1 to check for.
779  * \param[in]     pc2 Calculation 2 to check for.
780  * \param[in]     g1  Index group structure that contains the atoms from
781  *   \p pc1.
782  * \param[in,out] g   Working space, should have enough allocated memory to
783  *   contain the intersection of the atoms in \p pc1 and \p pc2.
784  * \returns   true if the two calculations should be merged to use a common
785  *   base, false otherwise.
786  */
787 static bool should_merge(gmx_ana_poscalc_t* pc1, gmx_ana_poscalc_t* pc2, gmx_ana_index_t* g1, gmx_ana_index_t* g)
788 {
789     gmx_ana_index_t g2;
790
791     /* Do not merge calculations with different mass weighting. */
792     if ((pc1->flags & POS_MASS) != (pc2->flags & POS_MASS))
793     {
794         return false;
795     }
796     /* Avoid messing up complete calculations. */
797     if ((pc1->flags & POS_COMPLWHOLE) != (pc2->flags & POS_COMPLWHOLE))
798     {
799         return false;
800     }
801     /* Find the overlap between the calculations. */
802     gmx_ana_index_set(&g2, pc2->b.nra, pc2->b.a, 0);
803     gmx_ana_index_intersection(g, g1, &g2);
804     /* Do not merge if there is no overlap. */
805     if (g->isize == 0)
806     {
807         return false;
808     }
809     /* Full completion calculations always match if the type is correct. */
810     if ((pc1->flags & POS_COMPLWHOLE) && (pc2->flags & POS_COMPLWHOLE) && pc1->type == pc2->type)
811     {
812         return true;
813     }
814     /* The calculations also match if the intersection consists of full
815      * blocks. */
816     if (gmx_ana_index_has_full_ablocks(g, &pc1->b) && gmx_ana_index_has_full_ablocks(g, &pc2->b))
817     {
818         return true;
819     }
820     return false;
821 }
822
823 /*! \brief
824  * Creates a static base for position calculation.
825  *
826  * \param     pc  Data structure to copy.
827  * \returns   Pointer to a newly allocated base for \p pc.
828  *
829  * Creates and returns a deep copy of \p pc, but clears the
830  * \ref POS_DYNAMIC and \ref POS_MASKONLY flags.
831  * The newly created structure is set as the base (\c gmx_ana_poscalc_t::sbase)
832  * of \p pc and inserted in the collection before \p pc.
833  */
834 static gmx_ana_poscalc_t* create_simple_base(gmx_ana_poscalc_t* pc)
835 {
836     gmx_ana_poscalc_t* base;
837     int                flags;
838
839     flags = pc->flags & ~(POS_DYNAMIC | POS_MASKONLY);
840     base  = pc->coll->createCalculation(pc->type, flags);
841     set_poscalc_maxindex(base, &pc->gmax, true);
842
843     base->p = new gmx_ana_pos_t();
844
845     pc->sbase = base;
846     pc->coll->removeCalculation(base);
847     pc->coll->insertCalculation(base, pc);
848
849     return base;
850 }
851
852 /*! \brief
853  * Merges a calculation into another calculation such that the new calculation
854  * can be used as a base.
855  *
856  * \param[in,out] base Base calculation to merge to.
857  * \param[in,out] pc   Position calculation to merge to \p base.
858  *
859  * After the call, \p base can be used as a base for \p pc (or any calculation
860  * that used it as a base).
861  * It is assumed that any overlap between \p base and \p pc is in complete
862  * blocks, i.e., that the merge is possible.
863  */
864 static void merge_to_base(gmx_ana_poscalc_t* base, gmx_ana_poscalc_t* pc)
865 {
866     gmx_ana_index_t gp, gb, g;
867     int             isize, bnr;
868     int             i, bi, bj, bo;
869
870     base->flags |= pc->flags & (POS_VELOCITIES | POS_FORCES);
871     gmx_ana_index_set(&gp, pc->b.nra, pc->b.a, 0);
872     gmx_ana_index_set(&gb, base->b.nra, base->b.a, 0);
873     isize = gmx_ana_index_difference_size(&gp, &gb);
874     if (isize > 0)
875     {
876         gmx_ana_index_clear(&g);
877         gmx_ana_index_reserve(&g, base->b.nra + isize);
878         /* Find the new blocks */
879         gmx_ana_index_difference(&g, &gp, &gb);
880         /* Count the blocks in g */
881         i = bi = bnr = 0;
882         while (i < g.isize)
883         {
884             while (pc->b.a[pc->b.index[bi]] != g.index[i])
885             {
886                 ++bi;
887             }
888             i += pc->b.index[bi + 1] - pc->b.index[bi];
889             ++bnr;
890             ++bi;
891         }
892         /* Merge the atoms into a temporary structure */
893         gmx_ana_index_merge(&g, &gb, &g);
894         /* Merge the blocks */
895         srenew(base->b.index, base->b.nr + bnr + 1);
896         i                     = g.isize - 1;
897         bi                    = base->b.nr - 1;
898         bj                    = pc->b.nr - 1;
899         bo                    = base->b.nr + bnr - 1;
900         base->b.index[bo + 1] = i + 1;
901         while (bo >= 0)
902         {
903             if (bi < 0 || base->b.a[base->b.index[bi + 1] - 1] != g.index[i])
904             {
905                 i -= pc->b.index[bj + 1] - pc->b.index[bj];
906                 --bj;
907             }
908             else
909             {
910                 if (bj >= 0 && pc->b.a[pc->b.index[bj + 1] - 1] == g.index[i])
911                 {
912                     --bj;
913                 }
914                 i -= base->b.index[bi + 1] - base->b.index[bi];
915                 --bi;
916             }
917             base->b.index[bo] = i + 1;
918             --bo;
919         }
920         base->b.nr += bnr;
921         base->b.nalloc_index += bnr;
922         sfree(base->b.a);
923         base->b.nra      = g.isize;
924         base->b.a        = g.index;
925         base->b.nalloc_a = g.isize;
926         /* Refresh the gmax field */
927         gmx_ana_index_set(&base->gmax, base->b.nra, base->b.a, 0);
928     }
929 }
930
931 /*! \brief
932  * Merges two bases into one.
933  *
934  * \param[in,out] tbase Base calculation to merge to.
935  * \param[in]     mbase Base calculation to merge to \p tbase.
936  *
937  * After the call, \p mbase has been freed and \p tbase is used as the base
938  * for all calculations that previously had \p mbase as their base.
939  * It is assumed that any overlap between \p tbase and \p mbase is in complete
940  * blocks, i.e., that the merge is possible.
941  */
942 static void merge_bases(gmx_ana_poscalc_t* tbase, gmx_ana_poscalc_t* mbase)
943 {
944     gmx_ana_poscalc_t* pc;
945
946     merge_to_base(tbase, mbase);
947     mbase->coll->removeCalculation(mbase);
948     /* Set tbase as the base for all calculations that had mbase */
949     pc = tbase->coll->first_;
950     while (pc)
951     {
952         if (pc->sbase == mbase)
953         {
954             pc->sbase = tbase;
955             tbase->refcount++;
956         }
957         pc = pc->next;
958     }
959     /* Free mbase */
960     mbase->refcount = 0;
961     gmx_ana_poscalc_free(mbase);
962 }
963
964 /*! \brief
965  * Setups the static base calculation for a position calculation.
966  *
967  * \param[in,out] pc   Position calculation to setup the base for.
968  */
969 static void setup_base(gmx_ana_poscalc_t* pc)
970 {
971     gmx_ana_poscalc_t *base, *pbase, *next;
972     gmx_ana_index_t    gp, g;
973
974     /* Exit immediately if pc should not have a base. */
975     if (!can_use_base(pc))
976     {
977         return;
978     }
979
980     gmx_ana_index_set(&gp, pc->b.nra, pc->b.a, 0);
981     gmx_ana_index_clear(&g);
982     gmx_ana_index_reserve(&g, pc->b.nra);
983     pbase = pc;
984     base  = pc->coll->first_;
985     while (base)
986     {
987         /* Save the next calculation so that we can safely delete base */
988         next = base->next;
989         /* Skip pc, calculations that already have a base (we should match the
990          * base instead), as well as calculations that should not have a base.
991          * If the above conditions are met, check whether we should do a
992          * merge.
993          */
994         if (base != pc && !base->sbase && can_use_base(base) && should_merge(pbase, base, &gp, &g))
995         {
996             /* Check whether this is the first base found */
997             if (pbase == pc)
998             {
999                 /* Create a real base if one is not present */
1000                 if (!base->p)
1001                 {
1002                     pbase = create_simple_base(base);
1003                 }
1004                 else
1005                 {
1006                     pbase = base;
1007                 }
1008                 /* Make it a base for pc as well */
1009                 merge_to_base(pbase, pc);
1010                 pc->sbase = pbase;
1011                 pbase->refcount++;
1012             }
1013             else /* This was not the first base */
1014             {
1015                 if (!base->p)
1016                 {
1017                     /* If it is not a real base, just make the new base as
1018                      * the base for it as well. */
1019                     merge_to_base(pbase, base);
1020                     base->sbase = pbase;
1021                     pbase->refcount++;
1022                 }
1023                 else
1024                 {
1025                     /* If base is a real base, merge it with the new base
1026                      * and delete it. */
1027                     merge_bases(pbase, base);
1028                 }
1029             }
1030             gmx_ana_index_set(&gp, pbase->b.nra, pbase->b.a, 0);
1031             gmx_ana_index_reserve(&g, pc->b.nra);
1032         }
1033         /* Proceed to the next unchecked calculation */
1034         base = next;
1035     }
1036
1037     gmx_ana_index_deinit(&g);
1038
1039     /* If no base was found, create one if one is required */
1040     if (!pc->sbase && (pc->flags & POS_DYNAMIC) && (pc->flags & (POS_COMPLMAX | POS_COMPLWHOLE)))
1041     {
1042         create_simple_base(pc);
1043     }
1044 }
1045
1046 /*!
1047  * \param[in,out] pc    Position calculation data structure.
1048  * \param[in]     flags New flags.
1049  *
1050  * \p flags are added to the old flags.
1051  * If calculation type is \ref POS_ATOM, \ref POS_MASS is automatically
1052  * cleared.
1053  * If both \ref POS_DYNAMIC and \ref POS_MASKONLY are provided,
1054  * \ref POS_DYNAMIC is cleared.
1055  * If calculation type is not \ref POS_RES or \ref POS_MOL,
1056  * \ref POS_COMPLMAX and \ref POS_COMPLWHOLE are automatically cleared.
1057  */
1058 void gmx_ana_poscalc_set_flags(gmx_ana_poscalc_t* pc, int flags)
1059 {
1060     if (pc->type == POS_ATOM)
1061     {
1062         flags &= ~POS_MASS;
1063     }
1064     if (flags & POS_MASKONLY)
1065     {
1066         flags &= ~POS_DYNAMIC;
1067     }
1068     if (pc->type != POS_RES && pc->type != POS_MOL)
1069     {
1070         flags &= ~(POS_COMPLMAX | POS_COMPLWHOLE);
1071     }
1072     pc->flags |= flags;
1073 }
1074
1075 /*!
1076  * \param[in,out] pc  Position calculation data structure.
1077  * \param[in]     g   Maximum index group for the calculation.
1078  *
1079  * Subsequent calls to gmx_ana_poscalc_update() should use only subsets of
1080  * \p g for evaluation.
1081  *
1082  * The topology should have been set for the collection of which \p pc is
1083  * a member.
1084  */
1085 void gmx_ana_poscalc_set_maxindex(gmx_ana_poscalc_t* pc, gmx_ana_index_t* g)
1086 {
1087     set_poscalc_maxindex(pc, g, false);
1088     setup_base(pc);
1089 }
1090
1091 /*!
1092  * \param[in]  pc  Position calculation data structure.
1093  * \param[out] p   Output positions.
1094  *
1095  * Calls to gmx_ana_poscalc_update() using \p pc should use only positions
1096  * initialized with this function.
1097  * The \c p->g pointer is initialized to point to an internal group that
1098  * contains the maximum index group set with gmx_ana_poscalc_set_maxindex().
1099  */
1100 void gmx_ana_poscalc_init_pos(gmx_ana_poscalc_t* pc, gmx_ana_pos_t* p)
1101 {
1102     gmx_ana_indexmap_init(&p->m, &pc->gmax, pc->coll->top_, pc->itype);
1103     /* Only do the static optimization when there is no completion */
1104     if (!(pc->flags & POS_DYNAMIC) && pc->b.nra == pc->gmax.isize)
1105     {
1106         gmx_ana_indexmap_set_static(&p->m, &pc->b);
1107     }
1108     gmx_ana_pos_reserve(p, p->m.mapb.nr, -1);
1109     if (pc->flags & POS_VELOCITIES)
1110     {
1111         gmx_ana_pos_reserve_velocities(p);
1112     }
1113     if (pc->flags & POS_FORCES)
1114     {
1115         gmx_ana_pos_reserve_forces(p);
1116     }
1117 }
1118
1119 /*!
1120  * \param  pc  Position calculation data to be freed.
1121  *
1122  * The \p pc pointer is invalid after the call.
1123  */
1124 void gmx_ana_poscalc_free(gmx_ana_poscalc_t* pc)
1125 {
1126     if (!pc)
1127     {
1128         return;
1129     }
1130
1131     pc->refcount--;
1132     if (pc->refcount > 0)
1133     {
1134         return;
1135     }
1136
1137     pc->coll->removeCalculation(pc);
1138     if (pc->b.nalloc_index > 0)
1139     {
1140         sfree(pc->b.index);
1141     }
1142     if (pc->b.nalloc_a > 0)
1143     {
1144         sfree(pc->b.a);
1145     }
1146     if (pc->flags & POS_COMPLWHOLE)
1147     {
1148         gmx_ana_index_deinit(&pc->gmax);
1149     }
1150     delete pc->p;
1151     if (pc->sbase)
1152     {
1153         gmx_ana_poscalc_free(pc->sbase);
1154         sfree(pc->baseid);
1155     }
1156     sfree(pc);
1157 }
1158
1159 gmx::PositionCalculationCollection::RequiredTopologyInfo gmx_ana_poscalc_required_topology_info(gmx_ana_poscalc_t* pc)
1160 {
1161     return gmx::requiredTopologyInfo(pc->type, pc->flags);
1162 }
1163
1164 /*!
1165  * \param[in]     pc   Position calculation data.
1166  * \param[in,out] p    Output positions, initialized previously with
1167  *   gmx_ana_poscalc_init_pos() using \p pc.
1168  * \param[in]     g    Index group to use for the update.
1169  * \param[in]     fr   Current frame.
1170  * \param[in]     pbc  PBC data, or NULL if no PBC should be used.
1171  *
1172  * gmx_ana_poscalc_init_frame() should be called for each frame before calling
1173  * this function.
1174  */
1175 void gmx_ana_poscalc_update(gmx_ana_poscalc_t* pc,
1176                             gmx_ana_pos_t*     p,
1177                             gmx_ana_index_t*   g,
1178                             t_trxframe*        fr,
1179                             const t_pbc*       pbc)
1180 {
1181     int i, bi, bj;
1182
1183     if (pc->bEval && !(pc->flags & POS_MASKONLY))
1184     {
1185         return;
1186     }
1187     if (pc->sbase)
1188     {
1189         gmx_ana_poscalc_update(pc->sbase, nullptr, nullptr, fr, pbc);
1190     }
1191     if (!p)
1192     {
1193         p = pc->p;
1194     }
1195     if (!g)
1196     {
1197         g = &pc->gmax;
1198     }
1199
1200     /* Update the index map */
1201     if (pc->flags & POS_DYNAMIC)
1202     {
1203         gmx_ana_indexmap_update(&p->m, g, false);
1204     }
1205     else if (pc->flags & POS_MASKONLY)
1206     {
1207         gmx_ana_indexmap_update(&p->m, g, true);
1208         if (pc->bEval)
1209         {
1210             return;
1211         }
1212     }
1213     if (!(pc->flags & POS_DYNAMIC))
1214     {
1215         pc->bEval = true;
1216     }
1217
1218     /* Evaluate the positions */
1219     if (pc->sbase)
1220     {
1221         /* TODO: It might be faster to evaluate the positions within this
1222          * loop instead of in the beginning. */
1223         if (pc->flags & POS_DYNAMIC)
1224         {
1225             for (bi = 0; bi < p->count(); ++bi)
1226             {
1227                 bj = pc->baseid[p->m.refid[bi]];
1228                 copy_rvec(pc->sbase->p->x[bj], p->x[bi]);
1229             }
1230             if (p->v)
1231             {
1232                 for (bi = 0; bi < p->count(); ++bi)
1233                 {
1234                     bj = pc->baseid[p->m.refid[bi]];
1235                     copy_rvec(pc->sbase->p->v[bj], p->v[bi]);
1236                 }
1237             }
1238             if (p->f)
1239             {
1240                 for (bi = 0; bi < p->count(); ++bi)
1241                 {
1242                     bj = pc->baseid[p->m.refid[bi]];
1243                     copy_rvec(pc->sbase->p->f[bj], p->f[bi]);
1244                 }
1245             }
1246         }
1247         else
1248         {
1249             for (bi = 0; bi < p->count(); ++bi)
1250             {
1251                 bj = pc->baseid[bi];
1252                 copy_rvec(pc->sbase->p->x[bj], p->x[bi]);
1253             }
1254             if (p->v)
1255             {
1256                 for (bi = 0; bi < p->count(); ++bi)
1257                 {
1258                     bj = pc->baseid[bi];
1259                     copy_rvec(pc->sbase->p->v[bj], p->v[bi]);
1260                 }
1261             }
1262             if (p->f)
1263             {
1264                 for (bi = 0; bi < p->count(); ++bi)
1265                 {
1266                     bj = pc->baseid[bi];
1267                     copy_rvec(pc->sbase->p->f[bj], p->f[bi]);
1268                 }
1269             }
1270         }
1271     }
1272     else /* pc->sbase is NULL */
1273     {
1274         if (pc->flags & POS_DYNAMIC)
1275         {
1276             pc->b.nr    = p->m.mapb.nr;
1277             pc->b.index = p->m.mapb.index;
1278             pc->b.nra   = g->isize;
1279             pc->b.a     = g->index;
1280         }
1281         if (p->v && !fr->bV)
1282         {
1283             for (i = 0; i < pc->b.nra; ++i)
1284             {
1285                 clear_rvec(p->v[i]);
1286             }
1287         }
1288         if (p->f && !fr->bF)
1289         {
1290             for (i = 0; i < pc->b.nra; ++i)
1291             {
1292                 clear_rvec(p->f[i]);
1293             }
1294         }
1295         gmx::ArrayRef<const int> index = pc->coll->getFrameIndices(pc->b.nra, pc->b.a);
1296         const gmx_mtop_t*        top   = pc->coll->top_;
1297         const bool               bMass = (pc->flags & POS_MASS) != 0;
1298         switch (pc->type)
1299         {
1300             case POS_ATOM:
1301                 for (i = 0; i < pc->b.nra; ++i)
1302                 {
1303                     copy_rvec(fr->x[index[i]], p->x[i]);
1304                 }
1305                 if (p->v && fr->bV)
1306                 {
1307                     for (i = 0; i < pc->b.nra; ++i)
1308                     {
1309                         copy_rvec(fr->v[index[i]], p->v[i]);
1310                     }
1311                 }
1312                 if (p->f && fr->bF)
1313                 {
1314                     for (i = 0; i < pc->b.nra; ++i)
1315                     {
1316                         copy_rvec(fr->f[index[i]], p->f[i]);
1317                     }
1318                 }
1319                 break;
1320             case POS_ALL:
1321                 gmx_calc_comg(top, fr->x, index.size(), index.data(), bMass, p->x[0]);
1322                 if (p->v && fr->bV)
1323                 {
1324                     gmx_calc_comg(top, fr->v, index.size(), index.data(), bMass, p->v[0]);
1325                 }
1326                 if (p->f && fr->bF)
1327                 {
1328                     gmx_calc_comg_f(top, fr->f, index.size(), index.data(), bMass, p->f[0]);
1329                 }
1330                 break;
1331             case POS_ALL_PBC:
1332                 gmx_calc_comg_pbc(top, fr->x, pbc, index.size(), index.data(), bMass, p->x[0]);
1333                 if (p->v && fr->bV)
1334                 {
1335                     gmx_calc_comg(top, fr->v, index.size(), index.data(), bMass, p->v[0]);
1336                 }
1337                 if (p->f && fr->bF)
1338                 {
1339                     gmx_calc_comg_f(top, fr->f, index.size(), index.data(), bMass, p->f[0]);
1340                 }
1341                 break;
1342             default:
1343                 // TODO: It would probably be better to do this without the type casts.
1344                 gmx_calc_comg_block(top, fr->x, reinterpret_cast<t_block*>(&pc->b), index.data(),
1345                                     bMass, p->x);
1346                 if (p->v && fr->bV)
1347                 {
1348                     gmx_calc_comg_block(top, fr->v, reinterpret_cast<t_block*>(&pc->b),
1349                                         index.data(), bMass, p->v);
1350                 }
1351                 if (p->f && fr->bF)
1352                 {
1353                     gmx_calc_comg_f_block(top, fr->f, reinterpret_cast<t_block*>(&pc->b),
1354                                           index.data(), bMass, p->f);
1355                 }
1356                 break;
1357         }
1358     }
1359 }