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