diff --git a/include/openmc/plot.h b/include/openmc/plot.h index 8a9e4e74d..63246b265 100644 --- a/include/openmc/plot.h +++ b/include/openmc/plot.h @@ -53,6 +53,7 @@ struct RGBColor { typedef xt::xtensor ImageData; typedef xt::xtensor IdData; +typedef xt::xtensor PropertyData; enum class PlotType { slice = 1, @@ -73,8 +74,17 @@ enum class PlotColorBy { //=============================================================================== // Plot class //=============================================================================== -struct PlotBase { +class PlotBase { + public: + IdData get_id_map() const; + PropertyData get_property_map() const; + + private: + template + D generate_data() const; + // Members + public: Position origin_; //!< Plot origin in geometry Position width_; //!< Plot width in geometry PlotBasis basis_; //!< Plot basis (XY/XZ/YZ) diff --git a/src/plot.cpp b/src/plot.cpp index 729b2db48..95f7acad7 100644 --- a/src/plot.cpp +++ b/src/plot.cpp @@ -29,6 +29,34 @@ namespace openmc { const RGBColor WHITE {255, 255, 255}; constexpr int PLOT_LEVEL_LOWEST {-1}; //!< lower bound on plot universe level constexpr int NOT_FOUND {-1}; + +struct id_setter { + void operator()(const Particle& p, IdData& ids, int y, int x, int level) { + Cell* c = model::cells[p.coord_[level].cell].get(); + ids(y,x,0) = c->id_; + if (c->type_ == FILL_UNIVERSE || p.material_ == MATERIAL_VOID) { + ids(y,x,1) = NOT_FOUND; + } else { + Material* m = model::materials[p.material_].get(); + ids(y,x,1) = m->id_; + } + } +}; + +struct property_setter { + void operator()(const Particle& p, PropertyData& props, int y, int x, int level) { + Cell* c = model::cells[p.coord_[level].cell].get(); + props(y,x,0) = (p.sqrtkT_ * p.sqrtkT_) * K_BOLTZMANN; + if (c->type_ == FILL_UNIVERSE || p.material_ == MATERIAL_VOID) { + props(y,x, 1) = NOT_FOUND; + } else { + Material* m = model::materials[p.material_].get(); + props(y,x,1) = m->density_gpcc_; + } + } +}; + + //============================================================================== // Global variables //============================================================================== @@ -637,6 +665,82 @@ Plot::Plot(pugi::xml_node plot_node) set_mask(plot_node); } // End Plot constructor +template +D PlotBase::generate_data() const { + + size_t width = pixels_[0]; + size_t height = pixels_[1]; + + // get pixel size + double in_pixel = (width_[0])/static_cast(width); + double out_pixel = (width_[1])/static_cast(height); + + // size data array + D data({height, width, 2}, NOT_FOUND); + + // setup basis indices and initial position centered on pixel + int in_i, out_i; + Position xyz = origin_; + switch(basis_) { + case PlotBasis::xy : + in_i = 0; + out_i = 1; + break; + case PlotBasis::xz : + in_i = 0; + out_i = 2; + break; + case PlotBasis::yz : + in_i = 1; + out_i = 2; + break; + } + + // set initial position + xyz[in_i] = origin_[in_i] - width_[0] / 2. + in_pixel / 2.; + xyz[out_i] = origin_[out_i] + width_[1] / 2. - out_pixel / 2.; + + // arbitrary direction + Direction dir = {0.5, 0.5, 0.5}; + +#pragma omp parallel +{ + Particle p; + p.r() = xyz; + p.u() = dir; + p.coord_[0].universe = model::root_universe; + int level = level_; + int j{}; + + #pragma omp for + for (int y = 0; y < height; y++) { + p.r()[out_i] = xyz[out_i] - out_pixel * y; + for (int x = 0; x < width; x++) { + p.r()[in_i] = xyz[in_i] + in_pixel * x; + p.n_coord_ = 1; + // local variables + bool found_cell = find_cell(&p, 0); + j = p.n_coord_ - 1; + if (level >=0) {j = level + 1;} + if (found_cell) { + setter()(p, data, y, x, j); + Cell* c = model::cells[p.coord_[j].cell].get(); + } + } // inner for + } // outer for + } // omp parallel + + return data; +} + +IdData PlotBase::get_id_map() const { + return generate_data(); +} + +PropertyData PlotBase::get_property_map() const { + return generate_data(); +} + //============================================================================== // POSITION_RGB computes the red/green/blue values for a given plot with the // current particle's position @@ -963,73 +1067,10 @@ extern "C" int openmc_id_map(const void* plot, int32_t* data_out) return OPENMC_E_INVALID_ARGUMENT; } - size_t width = plt->pixels_[0]; - size_t height = plt->pixels_[1]; - - // get pixel size - double in_pixel = (plt->width_[0])/static_cast(width); - double out_pixel = (plt->width_[1])/static_cast(height); - - // size data array - IdData data({height, width, 2}, NOT_FOUND); - - // setup basis indices and initial position centered on pixel - int in_i, out_i; - Position xyz = plt->origin_; - switch(plt->basis_) { - case PlotBasis::xy : - in_i = 0; - out_i = 1; - break; - case PlotBasis::xz : - in_i = 0; - out_i = 2; - break; - case PlotBasis::yz : - in_i = 1; - out_i = 2; - break; - } - - // set initial position - xyz[in_i] = plt->origin_[in_i] - plt->width_[0] / 2. + in_pixel / 2.; - xyz[out_i] = plt->origin_[out_i] + plt->width_[1] / 2. - out_pixel / 2.; - - // arbitrary direction - Direction dir = {0.5, 0.5, 0.5}; - - #pragma omp parallel - { - Particle p; - p.r() = xyz; - p.u() = dir; - p.coord_[0].universe = model::root_universe; - int level = plt->level_; - int j{}; - - #pragma omp for - for (int y = 0; y < height; y++) { - p.r()[out_i] = xyz[out_i] - out_pixel * y; - for (int x = 0; x < width; x++) { - p.r()[in_i] = xyz[in_i] + in_pixel * x; - p.n_coord_ = 1; - // local variables - bool found_cell = find_cell(&p, 0); - j = p.n_coord_ - 1; - if (level >=0) {j = level + 1;} - if (found_cell) { - Cell* c = model::cells[p.coord_[j].cell].get(); - data(y,x,0) = c->id_; - if (c->type_ != FILL_UNIVERSE && p.material_ != MATERIAL_VOID) { - data(y,x,1) = model::materials[p.material_]->id_; - } - } - } // inner for - } // outer for - } // omp parallel + auto ids = plt->get_id_map(); // write id data to array - std::copy(data.begin(), data.end(), data_out); + std::copy(ids.begin(), ids.end(), data_out); return 0; }