Separating Plot class into a C-compatible version for external use. Adding a function for returning an ID map.

include/openmc/plot.h

@@ -20,6 +20,7 @@ namespace openmc {
 //===============================================================================
 
 class Plot;
+class PlotC;
 
 namespace model {
 
@@ -52,6 +53,7 @@ struct RGBColor {
 };
 
 typedef xt::xtensor<RGBColor, 2> ImageData;
+typedef xt::xtensor<int32_t, 3>  IDData;
 
 enum class PlotType {
   slice = 1,
@@ -73,7 +75,27 @@ enum class PlotColorBy {
 // Plot class
 //===============================================================================
 
-class Plot
+
+// plot class for interaction with external code
+class CPlot {
+
+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
+};
+
+
+class Plot : public CPlot
 {
 
 public:
@@ -95,20 +117,8 @@ 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
   std::vector<RGBColor> colors_; //!< Plot colors
   std::string path_plot_; //!< Plot output filename
 };

openmc/capi/plot.py

@@ -1,15 +1,10 @@
-from ctypes import c_int, c_double, c_ushort, POINTER, Structure
+from ctypes import c_int, c_int32, c_double, c_ushort, POINTER, Structure
 
 from . import _dll
 from .core import _DLLGlobal
 
 import numpy as np
 
-_dll.openmc_get_image_data.argtypes = [c_int, c_int, c_int, c_int, c_int,
-                                       POINTER(c_double*3), c_double,
-                                       c_double, POINTER(c_double)]
-_dll.openmc_get_image_data.restype = c_double
-
 class _Position(Structure):
     pass
 
@@ -38,25 +33,15 @@ _Plot._fields_ = [('id_', c_int),
                   ('level_', c_int),
                   ('index_meshlines_mesh_', c_int),
                   ('meshlines_color_', _RGBColor),
-                  ('not_found_', _RGBColor),
-                  ('colors_', POINTER(_RGBColor))]
+                  ('not_found_', _RGBColor)]
 
 
-_dll.image_for_plot.argtypes= [POINTER(_Plot),]
-_dll.image_for_plot.restype = c_int
-
-def get_image_data():
-    origin = np.array([0.0, 0.0, 0.0])
-
-    val = np.array(1.0)
-
-    out = _dll.openmc_get_image_data(0, 0, 0, 0, 0, origin.ctypes.data_as(POINTER(c_double*3)), 0.0, 0.0, val.ctypes.data_as(POINTER(c_double)))
-    print(val)
-    return out
-
+_dll.openmc_id_map.argtypes= [POINTER(_Plot),]
+_dll.openmc_id_map.restype = c_int
 
 def image_data_for_plot(plot):
-    img_data = np.zeros((plot.pixels_[0], plot.pixels_[1], 3), dtype = float)
-
-    out = _dll.image_for_plot(POINTER(_Plot)(plot), img_data.ctypes.data_as(POINTER(c_double)))
+    img_data = np.zeros((plot.pixels_[0], plot.pixels_[1], 2), dtype=np.dtype('int32'))
+    print(img_data[0,0])
+    out = _dll.openmc_id_map(POINTER(_Plot)(plot), img_data.ctypes.data_as(POINTER(c_int32)))
+    print(img_data[0,0])
     return img_data

src/plot.cpp

@@ -645,9 +645,9 @@ 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);
   set_output_path(plot_node);
@@ -979,4 +979,85 @@ RGBColor random_color() {
   return {int(prn()*255), int(prn()*255), int(prn()*255)};
 }
 
+extern "C" int openmc_id_map(const CPlot& pl, int* data_out) {
+
+  size_t width = pl.pixels_[0];
+  size_t height = pl.pixels_[1];
+
+  double in_pixel = (pl.width_[0])/static_cast<double>(width);
+  double out_pixel = (pl.width_[1])/static_cast<double>(height);
+
+  IDData data;
+  data.resize({width, height, 2});
+
+  int in_i, out_i;
+  double xyz[3];
+  switch(pl.basis_) {
+  case PlotBasis::xy :
+    in_i = 0;
+    out_i = 1;
+    xyz[0] = pl.origin_[0] - pl.width_[0] / 2.;
+    xyz[1] = pl.origin_[1] + pl.width_[1] / 2.;
+    xyz[2] = pl.origin_[2];
+    break;
+  case PlotBasis::xz :
+    in_i = 0;
+    out_i = 2;
+    xyz[0] = pl.origin_[0] - pl.width_[0] / 2.;
+    xyz[1] = pl.origin_[1];
+    xyz[2] = pl.origin_[2] + pl.width_[1] / 2.;
+    break;
+  case PlotBasis::yz :
+    in_i = 1;
+    out_i = 2;
+    xyz[0] = pl.origin_[0];
+    xyz[1] = pl.origin_[1] - pl.width_[0] / 2.;
+    xyz[2] = pl.origin_[2] + pl.width_[1] / 2.;
+    break;
+  }
+
+  double dir[3] = {0.5, 0.5, 0.5};
+
+  //#pragma omp parallel
+{
+  Particle p;
+  p.initialize();
+  std::copy(xyz, xyz+3, p.coord[0].xyz);
+  std::copy(dir, dir+3, p.coord[0].uvw);
+  p.coord[0].universe = model::root_universe;
+  int level = pl.level_;
+  int j{};
+
+  //#pragma omp for
+  for (int y = 0; y < height; y++) {
+    p.coord[0].xyz[out_i] = xyz[out_i] - out_pixel * y;
+    p.n_coord = 1;
+    for (int x = 0; x < width; x++) {
+      p.coord[0].xyz[in_i] = xyz[in_i] + in_pixel * x;
+      // local variables
+      bool found_cell = find_cell(&p, 0);
+      j = p.n_coord - 1;
+      if (level >=0) {j = level + 1;}
+      if (!found_cell) {
+        data(x,y,0) = -1;
+        data(x,y,1) = -1;
+      } else {
+        Cell *c = model::cells[p.coord[j].cell];
+        data(x,y,0) = c->id_;
+        if (c->type_ == FILL_UNIVERSE || p.material == MATERIAL_VOID) {
+          data(x,y,1) = -1;
+        } else {
+          data(x,y,1) = model::materials[p.material - 1]->id_;
+        }
+      }
+    }
+  }
+ }
+
+  size_t i = 0;
+  for (auto v : data) { data_out[i++] = v; }
+
+  return 0;
+}
+
 } // namespace openmc