Line data    Source code

       1              : #ifndef _MOLECULARDYNAMICS_LINKEDCELL_H_
       2              : #define _MOLECULARDYNAMICS_LINKEDCELL_H_
       3              : 
       4              : #include "simplemd/Molecule.h"
       5              : #include <Kokkos_Core.hpp>
       6              : #include <iostream>
       7              : 
       8              : namespace simplemd {
       9              : class LinkedCell;
      10              : }
      11              : 
      12              : /**
      13              :  * A linked cell of a SimpleMD simulation for iterating the corresponding molecules.
      14              :  * NOTE: Never hold on to an instance of this class, as they are short lived contextual objects (host vs. device memory).
      15              :  * Always obtain them through simplemd::MoleculeContainer::operator[](int cellIndex).
      16              :  */
      17              : class simplemd::LinkedCell {
      18              : public:
      19              :   KOKKOS_FUNCTION LinkedCell() : _moleculeData(nullptr), _linkedCellNumMolecules(nullptr), _cellIndex(0), _isGhostCell(false) {}
      20              :   KOKKOS_FUNCTION LinkedCell(const Kokkos::View<Molecule**, Kokkos::LayoutRight, Kokkos::SharedSpace>* moleculeData,
      21              :                              const Kokkos::View<size_t*, Kokkos::LayoutRight, Kokkos::SharedSpace>* nMolecules, unsigned int cellIndex, bool isGhostCell)
      22              :       : _moleculeData(moleculeData), _linkedCellNumMolecules(nMolecules), _cellIndex(cellIndex), _isGhostCell(isGhostCell) {}
      23              : 
      24              :   class Iterator {
      25              :     using iterator_category = std::bidirectional_iterator_tag;
      26              :     using difference_type = std::ptrdiff_t;
      27              :     using value_type = Molecule;
      28              :     using pointer = value_type*;
      29              :     using reference = value_type&;
      30              : 
      31              :   public:
      32     33639488 :     KOKKOS_FUNCTION Iterator(pointer ptr, unsigned int idx) : _myPtr(ptr), _idx(idx) {}
      33              : 
      34              :     KOKKOS_INLINE_FUNCTION reference operator*() const { return *_myPtr; }
      35              :     KOKKOS_INLINE_FUNCTION pointer operator->() { return _myPtr; }
      36         3880 :     KOKKOS_INLINE_FUNCTION Iterator& operator++() {
      37         3880 :       _myPtr++;
      38         3880 :       _idx++;
      39         3860 :       return *this;
      40              :     }
      41         3860 :     KOKKOS_INLINE_FUNCTION Iterator operator++(int) {
      42         3860 :       Iterator temp = *this;
      43         3860 :       ++(*this);
      44         3860 :       return temp;
      45              :     }
      46           40 :     KOKKOS_INLINE_FUNCTION Iterator& operator--() {
      47           40 :       _myPtr--;
      48           40 :       _idx--;
      49           16 :       return *this;
      50              :     }
      51           24 :     KOKKOS_INLINE_FUNCTION Iterator operator--(int) {
      52           24 :       Iterator temp = *this;
      53           24 :       --(*this);
      54           24 :       return temp;
      55              :     }
      56              : 
      57              :     KOKKOS_INLINE_FUNCTION friend bool operator==(const Iterator& a, const Iterator& b) { return a._myPtr == b._myPtr; }
      58      8991592 :     KOKKOS_INLINE_FUNCTION friend bool operator!=(const Iterator& a, const Iterator& b) { return a._myPtr != b._myPtr; }
      59              : 
      60           80 :     KOKKOS_INLINE_FUNCTION unsigned int getIndex() const { return _idx; }
      61              : 
      62              :   private:
      63              :     pointer _myPtr;
      64              :     unsigned int _idx;
      65              :   };
      66              : 
      67     33635640 :   KOKKOS_FUNCTION Iterator begin() const { return Iterator(getMolecule(0), 0); }
      68     67286672 :   KOKKOS_FUNCTION Iterator end() const { return Iterator(getMolecule(numMolecules()), numMolecules()); }
      69              : 
      70         1306 :   KOKKOS_FUNCTION void insert(Molecule& molecule) {
      71         1306 :     *getMolecule(numMolecules()) = molecule;
      72         1306 :     changeMoleculeCount(+1);
      73         1306 :   }
      74            4 :   KOKKOS_FUNCTION void remove(int moleculeIdx) {
      75            4 :     *getMolecule(moleculeIdx) = *getMolecule(numMolecules() - 1);
      76            4 :     changeMoleculeCount(-1);
      77            4 :   }
      78       134132 :   KOKKOS_FUNCTION void clear() { (*_linkedCellNumMolecules)(_cellIndex) = 0; }
      79              : 
      80              :   std::string to_string() const {
      81              :     std::stringstream to_ret;
      82              :     to_ret << " numMol: " << numMolecules() << std::endl;
      83              :     return to_ret.str();
      84              :   }
      85              : 
      86     16822999 :   KOKKOS_INLINE_FUNCTION unsigned int numMolecules() const { return (*_linkedCellNumMolecules)(_cellIndex); }
      87              : 
      88              :   KOKKOS_INLINE_FUNCTION bool isGhostCell() const { return _isGhostCell; }
      89              : 
      90            0 :   KOKKOS_INLINE_FUNCTION size_t getIndex() const { return _cellIndex; }
      91              : 
      92              : private:
      93         1310 :   KOKKOS_INLINE_FUNCTION void changeMoleculeCount(int by) { (*_linkedCellNumMolecules)(_cellIndex) += by; }
      94              : 
      95     33640798 :   KOKKOS_INLINE_FUNCTION Molecule* getMolecule(unsigned int moleculeIndex) const { return &(*_moleculeData)(_cellIndex, moleculeIndex); }
      96              : 
      97              :   const Kokkos::View<Molecule**, Kokkos::LayoutRight, Kokkos::SharedSpace>* _moleculeData;
      98              :   const Kokkos::View<size_t*, Kokkos::LayoutRight, Kokkos::SharedSpace>* _linkedCellNumMolecules;
      99              :   const unsigned int _cellIndex;
     100              :   const bool _isGhostCell;
     101              : };
     102              : 
     103              : #endif // _MOLECULARDYNAMICS_LINKEDCELL_H_