diff --git a/src/lattice.cpp b/src/lattice.cpp index 93c53d88a..024e4cb48 100644 --- a/src/lattice.cpp +++ b/src/lattice.cpp @@ -49,6 +49,55 @@ Lattice::Lattice(pugi::xml_node lat_node) } } +//============================================================================== + +LatticeIter +Lattice::begin() +{ + return LatticeIter(*this, 0); +} + +LatticeIter +Lattice::end() +{ + return LatticeIter(*this, universes.size()); +} + +//============================================================================== + +void +Lattice::adjust_indices() +{ + // Adjust the indices for the universes array. + for (auto it = begin(); it != end(); ++it) { + int uid = *it; + auto search = universe_dict.find(uid); + if (search != universe_dict.end()) { + *it = search->second; + } else { + std::stringstream err_msg; + err_msg << "Invalid universe number " << uid << " specified on " + "lattice " << id; + fatal_error(err_msg); + } + } + + // Adjust the index for the outer universe. + if (outer != NO_OUTER_UNIVERSE) { + auto search = universe_dict.find(outer); + if (search != universe_dict.end()) { + outer = search->second; + } else { + std::stringstream err_msg; + err_msg << "Invalid universe number " << outer << " specified on " + "lattice " << id; + fatal_error(err_msg); + } + } +} + +//============================================================================== + void Lattice::to_hdf5(hid_t lattices_group) const { @@ -163,39 +212,6 @@ RectLattice::operator[](const int i_xyz[3]) //============================================================================== -void -RectLattice::adjust_indices() -{ - // Adjust the indices for the universes array. - for (auto it = this->begin(); it != this->end(); ++it) { - int uid = *it; - auto search = universe_dict.find(uid); - if (search != universe_dict.end()) { - *it = search->second; - } else { - std::stringstream err_msg; - err_msg << "Invalid universe number " << uid << " specified on " - "lattice " << id; - fatal_error(err_msg); - } - } - - // Adjust the index for the outer universe. - if (outer != NO_OUTER_UNIVERSE) { - auto search = universe_dict.find(outer); - if (search != universe_dict.end()) { - outer = search->second; - } else { - std::stringstream err_msg; - err_msg << "Invalid universe number " << outer << " specified on " - "lattice " << id; - fatal_error(err_msg); - } - } -} - -//============================================================================== - bool RectLattice::are_valid_indices(const int i_xyz[3]) const { @@ -521,49 +537,10 @@ HexLattice::operator[](const int i_xyz[3]) //============================================================================== -HexLatticeIter +LatticeIter HexLattice::begin() { - return HexLatticeIter(*this, n_rings-1, 0, 0); -} - -HexLatticeIter -HexLattice::end() -{ - return HexLatticeIter(*this, n_rings, 2*n_rings-2, n_axial-1); -} - -//============================================================================== - -void -HexLattice::adjust_indices() -{ - // Adjust the indices for the universes array. - for (auto it = this->begin(); it != this->end(); ++it) { - int uid = *it; - auto search = universe_dict.find(uid); - if (search != universe_dict.end()) { - *it = search->second; - } else { - std::stringstream err_msg; - err_msg << "Invalid universe number " << uid << " specified on lattice " - << id; - fatal_error(err_msg); - } - } - - // Adjust the index for the outer universe. - if (outer != NO_OUTER_UNIVERSE) { - auto search = universe_dict.find(outer); - if (search != universe_dict.end()) { - outer = search->second; - } else { - std::stringstream err_msg; - err_msg << "Invalid universe number " << outer << " specified on " - "lattice " << id; - fatal_error(err_msg); - } - } + return LatticeIter(*this, n_rings-1); } //============================================================================== @@ -771,6 +748,21 @@ HexLattice::get_local_xyz(const double global_xyz[3], const int i_xyz[3]) const //============================================================================== +bool +HexLattice::is_valid_index(int indx) const +{ + int nx {2*n_rings - 1}; + int ny {2*n_rings - 1}; + int nz {n_axial}; + int iz = indx / (nx * ny); + int iy = (indx - nx*ny*iz) / nx; + int ix = indx - nx*ny*iz - nx*iy; + int i_xyz[3] {ix+1, iy+1, iz+1}; // TODO: fix this off-by-one + return are_valid_indices(i_xyz); +} + +//============================================================================== + void HexLattice::to_hdf5_inner(hid_t lat_group) const { @@ -815,44 +807,6 @@ HexLattice::to_hdf5_inner(hid_t lat_group) const write_int(lat_group, 3, dims, "universes", out, false); } -//============================================================================== -// HexLatticeIter implementation -//============================================================================== - -HexLatticeIter& -HexLatticeIter::operator++() -{ - while (iz < nz) { - if (iy >= ny) { - // This axial layer is finished. Move to the next one. - ++iz; - iy = 0; - ix = n_rings - 2; - } else { - ++ix; - if (ix + iy < n_rings - 1) { - // We're in the lower-left dead zone. Keep iterating ix. - } else if (ix >= nx) { - // End of this row. Move to the next. - ix = -1; - ++iy; - } else if (ix + iy > 3*n_rings - 3) { - // We're in the upper-right dead zone. Move to the next row. - ix = -1; - ++iy; - } else { - return *this; - } - } - } - - // We've passed the end of the lattice. Return an end iterator. - ix = n_rings; - iy = 2*n_rings-2; - iz = nz-1; - return *this; -} - //============================================================================== // Non-method functions //============================================================================== diff --git a/src/lattice.h b/src/lattice.h index da07df511..938294ae4 100644 --- a/src/lattice.h +++ b/src/lattice.h @@ -36,6 +36,8 @@ extern std::map lattice_dict; //! Abstract type for ordered array of universes //============================================================================== +class LatticeIter; + class Lattice { public: @@ -51,8 +53,11 @@ public: virtual int32_t& operator[](const int i_xyz[3]) = 0; + virtual LatticeIter begin(); + LatticeIter end(); + //! Convert internal universe values from IDs to indices using universe_dict. - virtual void adjust_indices() = 0; + void adjust_indices(); //! Check lattice indices. //! @param i_xyz[3] The indices for a lattice tile. @@ -82,6 +87,15 @@ public: virtual std::array get_local_xyz(const double global_xyz[3], const int i_xyz[3]) const = 0; + //! Check flattened lattice index. + //! @param indx The index for a lattice tile. + //! @return true if the given index fit within the lattice bounds. False + //! otherwise. + virtual bool is_valid_index(int indx) const + { + return (indx > 0) && (indx < universes.size()); + } + //! Write all information needed to reconstruct the lattice to an HDF5 group. //! @param group_id An HDF5 group id. void to_hdf5(hid_t group_id) const; @@ -92,6 +106,44 @@ protected: virtual void to_hdf5_inner(hid_t group_id) const = 0; }; +class LatticeIter +{ +public: + LatticeIter(Lattice &lat_, int indx_) + : lat(lat_), + indx(indx_) + {} + + bool operator==(const LatticeIter &rhs) + { + return (&lat == &rhs.lat) && (indx == rhs.indx); + } + + bool operator!=(const LatticeIter &rhs) + { + return !(*this == rhs); + } + + int32_t& operator*() + { + return lat.universes[indx]; + } + + LatticeIter& operator++() + { + while (indx < lat.universes.size()) { + ++indx; + if (lat.is_valid_index(indx)) return *this; + } + indx = lat.universes.size(); + return *this; + } + + int indx; +private: + Lattice ⪫ +}; + //============================================================================== //============================================================================== @@ -104,11 +156,6 @@ public: int32_t& operator[](const int i_xyz[3]); - std::vector::iterator begin() {return universes.begin();} - std::vector::iterator end() {return universes.end();} - - void adjust_indices(); - bool are_valid_indices(const int i_xyz[3]) const; std::pair> @@ -130,8 +177,6 @@ protected: //============================================================================== //============================================================================== -class HexLatticeIter; - class HexLattice : public Lattice { public: @@ -141,10 +186,7 @@ public: int32_t& operator[](const int i_xyz[3]); - HexLatticeIter begin(); - HexLatticeIter end(); - - void adjust_indices(); + LatticeIter begin(); bool are_valid_indices(const int i_xyz[3]) const; @@ -156,6 +198,8 @@ public: std::array get_local_xyz(const double global_xyz[3], const int i_xyz[3]) const; + bool is_valid_index(int indx) const; + void to_hdf5_inner(hid_t group_id) const; protected: @@ -163,53 +207,6 @@ protected: int n_axial; //! Number of axial tile positions std::array center; //! Global center of lattice std::array pitch; //! Lattice tile width and height - - friend class HexLatticeIter; -}; - -class HexLatticeIter -{ -public: - HexLatticeIter(HexLattice &lat_, int ix_, int iy_, int iz_) - : lat(lat_), - n_rings(lat_.n_rings), - nx(2*lat_.n_rings-1), - ny(2*lat_.n_rings-1), - nz(lat_.n_axial), - ix(ix_), - iy(iy_), - iz(iz_) - {} - - bool operator==(const HexLatticeIter &rhs) - { - return (&lat == &rhs.lat) && (ix == rhs.ix) && (iy == rhs.iy) - && (iz == rhs.iz); - }; - - bool operator!=(const HexLatticeIter &rhs) - { - return !(*this == rhs); - }; - - int32_t& operator*() - { - int indx = nx*ny*iz + nx*iy + ix; - return lat.universes[indx]; - } - - HexLatticeIter& operator++(); - - HexLatticeIter operator++(int) - { - HexLatticeIter clone(*this); - ++(*this); - return clone; - } - -private: - HexLattice ⪫ - int n_rings, nx, ny, nz, ix, iy, iz; }; } // namespace openmc