diff --git a/include/openmc/capi.h b/include/openmc/capi.h index 9bfe4aa9ea..5d07edb212 100644 --- a/include/openmc/capi.h +++ b/include/openmc/capi.h @@ -70,6 +70,7 @@ extern "C" { int openmc_next_batch(int* status); int openmc_nuclide_name(int index, const char** name); int openmc_plot_geometry(); + int openmc_id_map(const void* slice, int32_t* data_out); int openmc_reset(); int openmc_run(); void openmc_set_seed(int64_t new_seed); diff --git a/include/openmc/plot.h b/include/openmc/plot.h index a6f215cfb4..8a9e4e74dc 100644 --- a/include/openmc/plot.h +++ b/include/openmc/plot.h @@ -52,6 +52,7 @@ struct RGBColor { }; typedef xt::xtensor ImageData; +typedef xt::xtensor IdData; enum class PlotType { slice = 1, @@ -72,8 +73,16 @@ enum class PlotColorBy { //=============================================================================== // Plot class //=============================================================================== +struct PlotBase { + // Members + Position origin_; //!< Plot origin in geometry + Position width_; //!< Plot width in geometry + PlotBasis basis_; //!< Plot basis (XY/XZ/YZ) + std::array pixels_; //!< Plot size in pixels + int level_; //!< Plot universe level +}; -class Plot +class Plot : public PlotBase { public: @@ -95,17 +104,12 @@ private: void set_meshlines(pugi::xml_node plot_node); void set_mask(pugi::xml_node plot_node); - // Members +// Members public: int id_; //!< Plot ID PlotType type_; //!< Plot type (Slice/Voxel) PlotColorBy color_by_; //!< Plot coloring (cell/material) - Position origin_; //!< Plot origin in geometry - Position width_; //!< Plot width in geometry - PlotBasis basis_; //!< Plot basis (XY/XZ/YZ) - std::array pixels_; //!< Plot size in pixels int meshlines_width_; //!< Width of lines added to the plot - int level_; //!< Plot universe level int index_meshlines_mesh_; //!< Index of the mesh to draw on the plot RGBColor meshlines_color_; //!< Color of meshlines on the plot RGBColor not_found_; //!< Plot background color diff --git a/openmc/capi/__init__.py b/openmc/capi/__init__.py index 1015df2166..5ca867d67e 100644 --- a/openmc/capi/__init__.py +++ b/openmc/capi/__init__.py @@ -42,7 +42,6 @@ else: def _dagmc_enabled(): return c_bool.in_dll(_dll, "dagmc_enabled").value - from .error import * from .core import * from .nuclide import * @@ -53,3 +52,4 @@ from .filter import * from .tally import * from .settings import settings from .math import * +from .plot import * diff --git a/openmc/capi/plot.py b/openmc/capi/plot.py new file mode 100644 index 0000000000..78fcef912c --- /dev/null +++ b/openmc/capi/plot.py @@ -0,0 +1,220 @@ +from ctypes import c_int, c_int32, c_double, Structure, POINTER + +from . import _dll +from .core import _DLLGlobal +from .error import _error_handler + +import numpy as np + + +class _Position(Structure): + """Definition of an xyz location in space with underlying c-types + + C-type Attributes + ----------------- + x : c_double + Position's x value (default: 0.0) + y : c_double + Position's y value (default: 0.0) + z : c_double + Position's z value (default: 0.0) + """ + _fields_ = [('x', c_double), + ('y', c_double), + ('z', c_double)] + + def __getitem__(self, idx): + if idx == 0: + return self.x + elif idx == 1: + return self.y + elif idx == 2: + return self.z + else: + raise IndexError("{} index is invalid for _Position".format(key)) + + def __setitem__(self, idx, val): + if idx == 0: + self.x = val + elif idx == 1: + self.y = val + elif idx == 2: + self.z = val + else: + raise IndexError("{} index is invalid for _Position".format(idx)) + + def __repr__(self): + return "({}, {}, {})".format(self.x, self.y, self.z) + + +class _PlotBase(Structure): + """A structure defining a 2-D geometry slice with underlying c-types + + C-Type Attributes + ----------------- + origin : openmc.capi.plot._Position + A position defining the origin of the plot. + width_ : openmc.capi.plot._Position + The width of the plot along the x, y, and z axes, respectively + basis_ : c_int + The axes basis of the plot view. + pixels_ : c_int[3] + The resolution of the plot in the horizontal and vertical dimensions + level_ : c_int + The universe level for the plot view + + Attributes + ---------- + origin : tuple or list of ndarray + Origin (center) of the plot + width : float + The horizontal dimension of the plot in geometry units (cm) + height : float + The vertical dimension of the plot in geometry units (cm) + basis : string + One of {'xy', 'xz', 'yz'} indicating the horizontal and vertical + axes of the plot. + h_res : float + The horizontal resolution of the plot in pixels + v_res : float + The vertical resolution of the plot in pixels + level : int + The universe level for the plot (default: -1 -> all universes shown) + """ + _fields_ = [('origin_', _Position), + ('width_', _Position), + ('basis_', c_int), + ('pixels_', 3*c_int), + ('level_', c_int)] + + def __init__(self): + self.level_ = -1 + + @property + def origin(self): + return self.origin_ + + @property + def width(self): + return self.width_.x + + @property + def height(self): + return self.width_.y + + @property + def basis(self): + if self.basis_ == 1: + return 'xy' + elif self.basis_ == 2: + return 'xz' + elif self.basis_ == 3: + return 'yz' + + raise ValueError("Plot basis {} is invalid".format(self.basis_)) + + @property + def h_res(self): + return self.pixels_[0] + + @property + def v_res(self): + return self.pixels_[1] + + @property + def level(self): + return int(self.level_) + + @origin.setter + def origin(self, origin): + self.origin_.x = origin[0] + self.origin_.y = origin[1] + self.origin_.z = origin[2] + + @width.setter + def width(self, width): + self.width_.x = width + + @height.setter + def height(self, height): + self.width_.y = height + + @basis.setter + def basis(self, basis): + if isinstance(basis, str): + valid_bases = ('xy', 'xz', 'yz') + basis = basis.lower() + if basis not in valid_bases: + raise ValueError("{} is not a valid plot basis.".format(basis)) + + if basis == 'xy': + self.basis_ = 1 + elif basis == 'xz': + self.basis_ = 2 + elif basis == 'yz': + self.basis_ = 3 + return + + if isinstance(basis, int): + valid_bases = (1, 2, 3) + if basis not in valid_bases: + raise ValueError("{} is not a valid plot basis.".format(basis)) + self.basis_ = basis + return + + raise ValueError("{} of type {} is an" + " invalid plot basis".format(basis, type(basis))) + + @h_res.setter + def h_res(self, h_res): + self.pixels_[0] = h_res + + @v_res.setter + def v_res(self, v_res): + self.pixels_[1] = v_res + + @level.setter + def level(self, level): + self.level_ = level + + def __repr__(self): + out_str = ["-----", + "Plot:", + "-----", + "Origin: {}".format(self.origin), + "Width: {}".format(self.width), + "Height: {}".format(self.height), + "Basis: {}".format(self.basis), + "HRes: {}".format(self.h_res), + "VRes: {}".format(self.v_res), + "Level: {}".format(self.level)] + return '\n'.join(out_str) + + +_dll.openmc_id_map.argtypes = [POINTER(_PlotBase), POINTER(c_int32)] +_dll.openmc_id_map.restype = c_int +_dll.openmc_id_map.errcheck = _error_handler + + +def id_map(plot): + """ + Generate a 2-D map of (cell_id, material_id). Used for in-memory image + generation. + + Parameters + ---------- + plot : openmc.capi.plot._PlotBase + Object describing the slice of the model to be generated + + Returns + ------- + id_map : numpy.ndarray + A NumPy array with shape (vertical pixels, horizontal pixels, 2) of + OpenMC property ids with dtype int32 + + """ + img_data = np.zeros((plot.v_res, plot.h_res, 2), + dtype=np.dtype('int32')) + _dll.openmc_id_map(POINTER(_PlotBase)(plot), + img_data.ctypes.data_as(POINTER(c_int32))) + return img_data diff --git a/src/plot.cpp b/src/plot.cpp index ae2f000237..18103a1413 100644 --- a/src/plot.cpp +++ b/src/plot.cpp @@ -1,5 +1,6 @@ #include "openmc/plot.h" +#include #include #include @@ -27,7 +28,7 @@ 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}; //============================================================================== // Global variables //============================================================================== @@ -133,26 +134,27 @@ void create_ppm(Plot pl) Direction u {0.5, 0.5, 0.5}; -#pragma omp parallel -{ - Particle p; - p.r() = r; - p.u() = u; - p.coord_[0].universe = model::root_universe; + #pragma omp parallel + { + Particle p; + p.r() = r; + p.u() = u; + p.coord_[0].universe = model::root_universe; -#pragma omp for - for (int y = 0; y < height; y++) { - p.r()[out_i] = r[out_i] - out_pixel * y; - for (int x = 0; x < width; x++) { - // local variables - RGBColor rgb; - int id; - p.r()[in_i] = r[in_i] + in_pixel * x; - position_rgb(p, pl, rgb, id); - data(x,y) = rgb; + #pragma omp for + for (int y = 0; y < height; y++) { + p.r()[out_i] = r[out_i] - out_pixel * y; + for (int x = 0; x < width; x++) { + // local variables + RGBColor rgb; + int id; + p.r()[in_i] = r[in_i] + in_pixel * x; + position_rgb(p, pl, rgb, id); + data(x,y) = rgb; + } } } -} + // draw mesh lines if present if (pl.index_meshlines_mesh_ >= 0) {draw_mesh_lines(pl, data);} @@ -618,8 +620,8 @@ Plot::set_mask(pugi::xml_node plot_node) } } -Plot::Plot(pugi::xml_node plot_node): -index_meshlines_mesh_(-1) +Plot::Plot(pugi::xml_node plot_node) + : index_meshlines_mesh_{-1} { set_id(plot_node); set_type(plot_node); @@ -657,7 +659,7 @@ void position_rgb(Particle p, Plot pl, RGBColor& rgb, int& id) if (!found_cell) { // If no cell, revert to default color rgb = pl.not_found_; - id = -1; + id = NOT_FOUND; } else { if (PlotColorBy::mats == pl.color_by_) { // Assign color based on material @@ -665,11 +667,11 @@ void position_rgb(Particle p, Plot pl, RGBColor& rgb, int& id) if (c->type_ == FILL_UNIVERSE) { // If we stopped on a middle universe level, treat as if not found rgb = pl.not_found_; - id = -1; + id = NOT_FOUND; } else if (p.material_ == MATERIAL_VOID) { // By default, color void cells white rgb = WHITE; - id = -1; + id = NOT_FOUND; } else { rgb = pl.colors_[p.material_]; id = model::materials[p.material_]->id_; @@ -952,4 +954,84 @@ RGBColor random_color() { return {int(prn()*255), int(prn()*255), int(prn()*255)}; } +extern "C" int openmc_id_map(const void* plot, int32_t* data_out) +{ + + auto plt = reinterpret_cast(plot); + if (!plt) { + set_errmsg("Invalid slice pointer passed to openmc_id_map"); + 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 + + // write id data to array + std::copy(data.begin(), data.end(), data_out); + + return 0; +} + } // namespace openmc diff --git a/tests/unit_tests/test_capi.py b/tests/unit_tests/test_capi.py index 0ebc898b48..97d18531bd 100644 --- a/tests/unit_tests/test_capi.py +++ b/tests/unit_tests/test_capi.py @@ -400,3 +400,22 @@ def test_load_nuclide(capi_init): # load non-existent nuclide with pytest.raises(exc.DataError): openmc.capi.load_nuclide('Pu3') + + +def test_id_map(capi_init): + expected_ids = np.array([[(3, 3), (2, 2), (3, 3)], + [(2, 2), (1, 1), (2, 2)], + [(3, 3), (2, 2), (3, 3)]], dtype='int32') + + # create a plot object + s = openmc.capi.plot._PlotBase() + s.width = 1.26 + s.height = 1.26 + s.v_res = 3 + s.h_res = 3 + s.origin = (0.0, 0.0, 0.0) + s.basis = 'xy' + s.level = -1 + + ids = openmc.capi.plot.id_map(s) + assert np.array_equal(expected_ids, ids)