Merge pull request #1190 from pshriwise/capi_id_map

Add ID map function to the capi

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);

include/openmc/plot.h

@@ -52,6 +52,7 @@ struct RGBColor {
 };
 
 typedef xt::xtensor<RGBColor, 2> ImageData;
+typedef xt::xtensor<int32_t, 3>  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<int, 3> 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<int, 3> 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

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 *

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

src/plot.cpp

@@ -1,5 +1,6 @@
 #include "openmc/plot.h"
 
+#include <algorithm>
 #include <fstream>
 #include <sstream>
 
@@ -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<const PlotBase*>(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<double>(width);
+  double out_pixel = (plt->width_[1])/static_cast<double>(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

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)