From 01b023b7ec7a4ad391afd293d514c1352d1292e3 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Thu, 24 Jun 2021 12:09:59 +0700
Subject: [PATCH] Include cell instance in id_map

---
 openmc/lib/plot.py | 2 +-
 src/plot.cpp       | 8 +++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/openmc/lib/plot.py b/openmc/lib/plot.py
index c51000e0d..b029d6757 100644
--- a/openmc/lib/plot.py
+++ b/openmc/lib/plot.py
@@ -232,7 +232,7 @@ def id_map(plot):
         OpenMC property ids with dtype int32
 
     """
-    img_data = np.zeros((plot.v_res, plot.h_res, 2),
+    img_data = np.zeros((plot.v_res, plot.h_res, 3),
                         dtype=np.dtype('int32'))
     _dll.openmc_id_map(plot, img_data.ctypes.data_as(POINTER(c_int32)))
     return img_data
diff --git a/src/plot.cpp b/src/plot.cpp
index 615701341..3e55d1d5f 100644
--- a/src/plot.cpp
+++ b/src/plot.cpp
@@ -36,7 +36,7 @@ constexpr int32_t NOT_FOUND {-2};
 constexpr int32_t OVERLAP {-3};
 
 IdData::IdData(size_t h_res, size_t v_res)
-  : data_({v_res, h_res, 2}, NOT_FOUND)
+  : data_({v_res, h_res, 3}, NOT_FOUND)
 { }
 
 void
@@ -44,18 +44,20 @@ IdData::set_value(size_t y, size_t x, const Particle& p, int level) {
   // set cell data
   if (p.n_coord() <= level) {
     data_(y, x, 0) = NOT_FOUND;
+    data_(y, x, 1) = NOT_FOUND;
   } else {
     data_(y, x, 0) = model::cells.at(p.coord(level).cell)->id_;
+    data_(y, x, 1) = p.cell_instance();
   }
 
   // set material data
   Cell* c = model::cells.at(p.coord(p.n_coord() - 1).cell).get();
   if (p.material() == MATERIAL_VOID) {
-    data_(y, x, 1) = MATERIAL_VOID;
+    data_(y, x, 2) = MATERIAL_VOID;
     return;
   } else if (c->type_ == Fill::MATERIAL) {
     Material* m = model::materials.at(p.material()).get();
-    data_(y, x, 1) = m->id_;
+    data_(y, x, 2) = m->id_;
   }
 }