1 /*
2  * Atom iterators
3  *
4  * (c) 2014 Schrodinger, Inc.
5  */
6 
7 #ifndef _H_ATOMITERATORS
8 #define _H_ATOMITERATORS
9 
10 #include "PyMOLGlobals.h"
11 #include "ObjectMolecule.h"
12 #include "CoordSet.h"
13 #include "AtomInfo.h"
14 
15 /**
16  * Atom iterator base class
17  *
18  * The initial state of the iterator points "before" the first atom, there is no
19  * current atom until next() is called the first time.
20  */
21 class AbstractAtomIterator {
22 protected:
23   int atm;      //!< Current atom index in object molecule
24   int idx = -1; //!< Current coordinate index (atom index in coordset)
25 
26 public:
27   ObjectMolecule* obj;    //!< Current object molecule
28   CoordSet* cs = nullptr; //!< Current coordiante set
29 
30   // virtual destructor for dynamic types
31   virtual ~AbstractAtomIterator() = default;
32 
33   /// Resets the internal state to start over
34   virtual void reset() = 0;
35 
36   /// Advance the internal state to the next atom.
37   /// @return False if there is no next atom
38   virtual bool next() = 0;
39 
40   /// Current atom
getAtomInfo()41   AtomInfoType * getAtomInfo() {
42     return obj->AtomInfo + atm;
43   };
getAtomInfo()44   const AtomInfoType * getAtomInfo() const {
45     return obj->AtomInfo + atm;
46   };
47 
48   /// Current atom's coordinates
getCoord()49   float * getCoord() {
50     return cs->coordPtr(idx);
51   };
52 
53   /// Current atom index in object molecule
getAtm()54   int getAtm() const {
55     return atm;
56   }
57 
58   /// Current coordinate index (atom index in coordset)
getIdx()59   int getIdx() const {
60     return idx;
61   }
62 };
63 
64 /**
65  * State specific iterator over an atom selection. Similar to `cmd.iterate_state`.
66  * If state is -1 then iterate over all states.
67  *
68  * @verbatim
69    SeleCoordIterator iter(G, sele, state);
70    while(iter.next()) {
71      dump3f(iter.getCoord(), "coords");
72    }
73    @endverbatim
74  */
75 class SeleCoordIterator : public AbstractAtomIterator {
76   PyMOLGlobals* G = nullptr;
77   int statearg; // state argument, can be -1 (all), -2 (current), -3 (effective)
78   int statemax; // largest state in selection
79   bool per_object;              // whether to iterate over object states or global states
80   ObjectMolecule * prev_obj;    // for per_object=true
81   int sele = -1;
82 
83 public:
84   int a;        //!< index in selection
85   int state;    //!< current state
86 
87   /// Undefined state until copy assigned from a valid state
88   SeleCoordIterator() = default;
89   SeleCoordIterator(
90       PyMOLGlobals* G_, int sele_, int state_, bool update_table = true);
91 
92   void reset();
93   bool next();
94 
95   /// Return true if iterating over all states
isMultistate()96   bool isMultistate() const {
97     return statearg == -1;
98   }
99 
100   /// Return true if iterating over all states of an object before advancing
101   /// to the next object, rather than iterating over global states
isPerObject()102   bool isPerObject() const {
103     return per_object;
104   }
105 
106   /// @see isPerObject()
setPerObject(bool per_object_)107   void setPerObject(bool per_object_) {
108     per_object = per_object_ && isMultistate();
109   }
110 
111 private:
nextStateInPrevObject()112   bool nextStateInPrevObject() {
113     if (prev_obj && (++state) < prev_obj->NCSet) {
114       a = prev_obj->SeleBase - 1;
115       return true;
116     }
117     return false;
118   }
119 };
120 
121 /**
122  * Iterator over an atom selection. Similar to `cmd.iterate`.
123  *
124  * Does NOT provide coord or coordset access
125  *
126  * @verbatim
127    SeleAtomIterator iter(G, sele);
128    while(iter.next()) {
129      ai = iter.getAtomInfo();
130    }
131    @endverbatim
132  */
133 class SeleAtomIterator : public AbstractAtomIterator {
134   PyMOLGlobals * G;
135   int sele;     // selection
136   char * stmp;  // temporary named selection
137 
138 public:
139   int a;        //!< index in selection
140 
141   /// Iterate over an existing atom selection
142   /// @pre ::SelectorUpdateTable was called
SeleAtomIterator(PyMOLGlobals * G_,int sele_)143   SeleAtomIterator(PyMOLGlobals * G_, int sele_) {
144     G = G_;
145     sele = sele_;
146     stmp = NULL;
147     reset();
148   };
149 
150   SeleAtomIterator(PyMOLGlobals * G_, const char * sele_);
151   ~SeleAtomIterator();
152 
153   void reset();
154   bool next();
155 };
156 
157 /**
158  * CoordSet atom iterator, iterates over sorted atoms
159  *
160  * @verbatim
161    CoordSetAtomIterator iter(G, cs);
162    while(iter.next()) {
163      dump3f(iter.getCoord(), "coords");
164    }
165    @endverbatim
166  */
167 class CoordSetAtomIterator : public AbstractAtomIterator {
168 public:
169 
CoordSetAtomIterator(CoordSet * cs_)170   CoordSetAtomIterator(CoordSet * cs_) {
171     cs = cs_;
172     obj = cs->Obj;
173     reset();
174   }
175 
reset()176   void reset() {
177     atm = -1;
178   }
179 
180   bool next();
181 };
182 
183 
184 #endif
185 
186 // vi:sw=2:expandtab
187