Make lattice iteration polymorphic

This commit is contained in:
Sterling Harper 2018-05-11 01:28:38 -04:00
parent 366bd36107
commit a0fd4e703f
2 changed files with 122 additions and 171 deletions

View file

@ -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
//==============================================================================

View file

@ -36,6 +36,8 @@ extern std::map<int32_t, int32_t> 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<double, 3>
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 &lat;
};
//==============================================================================
//==============================================================================
@ -104,11 +156,6 @@ public:
int32_t& operator[](const int i_xyz[3]);
std::vector<int32_t>::iterator begin() {return universes.begin();}
std::vector<int32_t>::iterator end() {return universes.end();}
void adjust_indices();
bool are_valid_indices(const int i_xyz[3]) const;
std::pair<double, std::array<int, 3>>
@ -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<double, 3>
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<double, 3> center; //! Global center of lattice
std::array<double, 2> 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 &lat;
int n_rings, nx, ny, nz, ix, iy, iz;
};
} // namespace openmc