Local replacement for _PlotBase since it's removed in OpenMC

Author: paulromano

openmc_plotter/plotmodel.py

@@ -2,7 +2,7 @@
 from ast import literal_eval
 from collections import defaultdict
 import copy
-from ctypes import c_int32, c_char_p
+from ctypes import c_int32, c_char_p, Structure, c_int, c_size_t, c_bool, c_double
 from dataclasses import dataclass
 import hashlib
 import itertools
@@ -1072,7 +1072,238 @@ def mat_ids(self):
         return self.geom_data[:, :, 2]
 
 
-class ViewParam(openmc.lib.plot._PlotBase):
+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(f"{idx} index is invalid for _Position")
+
+    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(f"{idx} index is invalid for _Position")
+
+    def __repr__(self):
+        return f"({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.lib.plot._Position
+        A position defining the origin of the plot.
+    u_span_ : openmc.lib.plot._Position
+        Full-width span vector defining the plot's horizontal axis.
+    v_span_ : openmc.lib.plot._Position
+        Full-height span vector defining the plot's vertical axis.
+    width_ : openmc.lib.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_size_t[3]
+        The resolution of the plot in the horizontal and vertical dimensions
+    color_overlaps_ : c_bool
+        Whether to assign unique IDs (-3) to overlapping regions.
+    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 : int
+        The horizontal resolution of the plot in pixels
+    v_res : int
+        The vertical resolution of the plot in pixels
+    level : int
+        The universe level for the plot (default: -1 -> all universes shown)
+    """
+    _fields_ = [('origin_', _Position),
+                ('u_span_', _Position),
+                ('v_span_', _Position),
+                ('width_', _Position),
+                ('basis_', c_int),
+                ('pixels_', 3*c_size_t),
+                ('color_overlaps_', c_bool),
+                ('level_', c_int)]
+
+    def __init__(self):
+        self.level_ = -1
+        self.basis_ = 1
+        self.color_overlaps_ = False
+        self._update_spans()
+
+    def _update_spans(self):
+        if self.basis_ == 1:
+            self.u_span_.x = self.width_.x
+            self.u_span_.y = 0.0
+            self.u_span_.z = 0.0
+            self.v_span_.x = 0.0
+            self.v_span_.y = self.width_.y
+            self.v_span_.z = 0.0
+        elif self.basis_ == 2:
+            self.u_span_.x = self.width_.x
+            self.u_span_.y = 0.0
+            self.u_span_.z = 0.0
+            self.v_span_.x = 0.0
+            self.v_span_.y = 0.0
+            self.v_span_.z = self.width_.y
+        elif self.basis_ == 3:
+            self.u_span_.x = 0.0
+            self.u_span_.y = self.width_.x
+            self.u_span_.z = 0.0
+            self.v_span_.x = 0.0
+            self.v_span_.y = 0.0
+            self.v_span_.z = self.width_.y
+
+    @property
+    def origin(self):
+        return self.origin_
+
+    @origin.setter
+    def origin(self, origin):
+        self.origin_.x = origin[0]
+        self.origin_.y = origin[1]
+        self.origin_.z = origin[2]
+
+    @property
+    def width(self):
+        return self.width_.x
+
+    @width.setter
+    def width(self, width):
+        self.width_.x = width
+        self._update_spans()
+
+    @property
+    def height(self):
+        return self.width_.y
+
+    @height.setter
+    def height(self, height):
+        self.width_.y = height
+        self._update_spans()
+
+    @property
+    def basis(self):
+        if self.basis_ == 1:
+            return 'xy'
+        elif self.basis_ == 2:
+            return 'xz'
+        elif self.basis_ == 3:
+            return 'yz'
+
+        raise ValueError(f"Plot basis {self.basis_} is invalid")
+
+    @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(f"{basis} is not a valid plot basis.")
+
+            if basis == 'xy':
+                self.basis_ = 1
+            elif basis == 'xz':
+                self.basis_ = 2
+            elif basis == 'yz':
+                self.basis_ = 3
+            self._update_spans()
+            return
+
+        if isinstance(basis, int):
+            valid_bases = (1, 2, 3)
+            if basis not in valid_bases:
+                raise ValueError(f"{basis} is not a valid plot basis.")
+            self.basis_ = basis
+            self._update_spans()
+            return
+
+        raise ValueError(f"{basis} of type {type(basis)} is an invalid plot basis")
+
+    @property
+    def h_res(self):
+        return self.pixels_[0]
+
+    @h_res.setter
+    def h_res(self, h_res):
+        self.pixels_[0] = h_res
+
+    @property
+    def v_res(self):
+        return self.pixels_[1]
+
+    @v_res.setter
+    def v_res(self, v_res):
+        self.pixels_[1] = v_res
+
+    @property
+    def level(self):
+        return int(self.level_)
+
+    @level.setter
+    def level(self, level):
+        self.level_ = level
+
+    @property
+    def color_overlaps(self):
+        return self.color_overlaps_
+
+    @color_overlaps.setter
+    def color_overlaps(self, color_overlaps):
+        self.color_overlaps_ = color_overlaps
+
+    def __repr__(self):
+        out_str = ["-----",
+                   "Plot:",
+                   "-----",
+                   f"Origin: {self.origin}",
+                   f"Width: {self.width}",
+                   f"Height: {self.height}",
+                   f"Basis: {self.basis}",
+                   f"HRes: {self.h_res}",
+                   f"VRes: {self.v_res}",
+                   f"Color Overlaps: {self.color_overlaps}",
+                   f"Level: {self.level}"]
+        return '\n'.join(out_str)
+
+
+class ViewParam(_PlotBase):
     """Viewer settings that are needed for _PlotBase and are independent
     of all other plotter/model settings.
 
@@ -1167,6 +1398,7 @@ def urc(self):
     def __eq__(self, other):
         return repr(self) == repr(other)
 
+
 class PlotViewIndependent:
     """View settings for OpenMC plot, independent of the model.