Source code for openmc.lib.plot

from collections.abc import Mapping
from ctypes import (c_bool, c_int, c_size_t, c_int32,
                    c_double, c_uint8, Structure, POINTER)
from weakref import WeakValueDictionary

from ..exceptions import AllocationError, InvalidIDError
from . import _dll
from .core import _FortranObjectWithID
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(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.
    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),
                ('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

    @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

    @property
    def height(self):
        return self.width_.y

    @height.setter
    def height(self, height):
        self.width_.y = height

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


_dll.openmc_id_map.argtypes = [POINTER(_PlotBase), POINTER(c_int32)]
_dll.openmc_id_map.restype = c_int
_dll.openmc_id_map.errcheck = _error_handler


[docs] def id_map(plot): """ Generate a 2-D map of cell and material IDs. Used for in-memory image generation. Parameters ---------- plot : openmc.lib.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, 3) of OpenMC property ids with dtype int32. The last dimension of the array contains, in order, cell IDs, cell instances, and material IDs. """ 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
_dll.openmc_property_map.argtypes = [POINTER(_PlotBase), POINTER(c_double)] _dll.openmc_property_map.restype = c_int _dll.openmc_property_map.errcheck = _error_handler
[docs] def property_map(plot): """ Generate a 2-D map of cell temperatures and material densities. Used for in-memory image generation. Parameters ---------- plot : openmc.lib.plot._PlotBase Object describing the slice of the model to be generated Returns ------- property_map : numpy.ndarray A NumPy array with shape (vertical pixels, horizontal pixels, 2) of OpenMC property ids with dtype float """ prop_data = np.zeros((plot.v_res, plot.h_res, 2)) _dll.openmc_property_map(plot, prop_data.ctypes.data_as(POINTER(c_double))) return prop_data
_dll.openmc_get_plot_index.argtypes = [c_int32, POINTER(c_int32)] _dll.openmc_get_plot_index.restype = c_int _dll.openmc_get_plot_index.errcheck = _error_handler _dll.openmc_plot_get_id.argtypes = [c_int32, POINTER(c_int32)] _dll.openmc_plot_get_id.restype = c_int _dll.openmc_plot_get_id.errcheck = _error_handler _dll.openmc_plot_set_id.argtypes = [c_int32, c_int32] _dll.openmc_plot_set_id.restype = c_int _dll.openmc_plot_set_id.errcheck = _error_handler _dll.openmc_plots_size.restype = c_size_t _dll.openmc_solidraytrace_plot_create.argtypes = [POINTER(c_int32)] _dll.openmc_solidraytrace_plot_create.restype = c_int _dll.openmc_solidraytrace_plot_create.errcheck = _error_handler _dll.openmc_solidraytrace_plot_get_pixels.argtypes = [ c_int32, POINTER(c_int32), POINTER(c_int32)] _dll.openmc_solidraytrace_plot_get_pixels.restype = c_int _dll.openmc_solidraytrace_plot_get_pixels.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_pixels.argtypes = [c_int32, c_int32, c_int32] _dll.openmc_solidraytrace_plot_set_pixels.restype = c_int _dll.openmc_solidraytrace_plot_set_pixels.errcheck = _error_handler _dll.openmc_solidraytrace_plot_get_color_by.argtypes = [c_int32, POINTER(c_int32)] _dll.openmc_solidraytrace_plot_get_color_by.restype = c_int _dll.openmc_solidraytrace_plot_get_color_by.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_color_by.argtypes = [c_int32, c_int32] _dll.openmc_solidraytrace_plot_set_color_by.restype = c_int _dll.openmc_solidraytrace_plot_set_color_by.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_default_colors.argtypes = [c_int32] _dll.openmc_solidraytrace_plot_set_default_colors.restype = c_int _dll.openmc_solidraytrace_plot_set_default_colors.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_all_opaque.argtypes = [c_int32] _dll.openmc_solidraytrace_plot_set_all_opaque.restype = c_int _dll.openmc_solidraytrace_plot_set_all_opaque.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_opaque.argtypes = [c_int32, c_int32, c_bool] _dll.openmc_solidraytrace_plot_set_opaque.restype = c_int _dll.openmc_solidraytrace_plot_set_opaque.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_color.argtypes = [c_int32, c_int32, c_uint8, c_uint8, c_uint8] _dll.openmc_solidraytrace_plot_set_color.restype = c_int _dll.openmc_solidraytrace_plot_set_color.errcheck = _error_handler _dll.openmc_solidraytrace_plot_get_camera_position.argtypes = [ c_int32, POINTER(c_double), POINTER(c_double), POINTER(c_double)] _dll.openmc_solidraytrace_plot_get_camera_position.restype = c_int _dll.openmc_solidraytrace_plot_get_camera_position.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_camera_position.argtypes = [c_int32, c_double, c_double, c_double] _dll.openmc_solidraytrace_plot_set_camera_position.restype = c_int _dll.openmc_solidraytrace_plot_set_camera_position.errcheck = _error_handler _dll.openmc_solidraytrace_plot_get_look_at.argtypes = [ c_int32, POINTER(c_double), POINTER(c_double), POINTER(c_double)] _dll.openmc_solidraytrace_plot_get_look_at.restype = c_int _dll.openmc_solidraytrace_plot_get_look_at.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_look_at.argtypes = [c_int32, c_double, c_double, c_double] _dll.openmc_solidraytrace_plot_set_look_at.restype = c_int _dll.openmc_solidraytrace_plot_set_look_at.errcheck = _error_handler _dll.openmc_solidraytrace_plot_get_up.argtypes = [ c_int32, POINTER(c_double), POINTER(c_double), POINTER(c_double)] _dll.openmc_solidraytrace_plot_get_up.restype = c_int _dll.openmc_solidraytrace_plot_get_up.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_up.argtypes = [c_int32, c_double, c_double, c_double] _dll.openmc_solidraytrace_plot_set_up.restype = c_int _dll.openmc_solidraytrace_plot_set_up.errcheck = _error_handler _dll.openmc_solidraytrace_plot_get_light_position.argtypes = [ c_int32, POINTER(c_double), POINTER(c_double), POINTER(c_double)] _dll.openmc_solidraytrace_plot_get_light_position.restype = c_int _dll.openmc_solidraytrace_plot_get_light_position.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_light_position.argtypes = [c_int32, c_double, c_double, c_double] _dll.openmc_solidraytrace_plot_set_light_position.restype = c_int _dll.openmc_solidraytrace_plot_set_light_position.errcheck = _error_handler _dll.openmc_solidraytrace_plot_get_fov.argtypes = [c_int32, POINTER(c_double)] _dll.openmc_solidraytrace_plot_get_fov.restype = c_int _dll.openmc_solidraytrace_plot_get_fov.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_fov.argtypes = [c_int32, c_double] _dll.openmc_solidraytrace_plot_set_fov.restype = c_int _dll.openmc_solidraytrace_plot_set_fov.errcheck = _error_handler _dll.openmc_solidraytrace_plot_update_view.argtypes = [c_int32] _dll.openmc_solidraytrace_plot_update_view.restype = c_int _dll.openmc_solidraytrace_plot_update_view.errcheck = _error_handler _dll.openmc_solidraytrace_plot_create_image.argtypes = [c_int32, POINTER(c_uint8), c_int32, c_int32] _dll.openmc_solidraytrace_plot_create_image.restype = c_int _dll.openmc_solidraytrace_plot_create_image.errcheck = _error_handler _dll.openmc_solidraytrace_plot_get_color.argtypes = [c_int32, c_int32, POINTER(c_uint8), POINTER(c_uint8), POINTER(c_uint8)] _dll.openmc_solidraytrace_plot_get_color.restype = c_int _dll.openmc_solidraytrace_plot_get_color.errcheck = _error_handler _dll.openmc_solidraytrace_plot_get_diffuse_fraction.argtypes = [ c_int32, POINTER(c_double)] _dll.openmc_solidraytrace_plot_get_diffuse_fraction.restype = c_int _dll.openmc_solidraytrace_plot_get_diffuse_fraction.errcheck = _error_handler _dll.openmc_solidraytrace_plot_set_diffuse_fraction.argtypes = [c_int32, c_double] _dll.openmc_solidraytrace_plot_set_diffuse_fraction.restype = c_int _dll.openmc_solidraytrace_plot_set_diffuse_fraction.errcheck = _error_handler
[docs] class SolidRayTracePlot(_FortranObjectWithID): """Solid ray-traced plot stored internally. This class exposes a solid ray-traced plot that is stored internally in the OpenMC library. To obtain a view of an existing plot with a given ID, use the :data:`openmc.lib.plots` mapping. Parameters ---------- uid : int or None Unique ID of the plot new : bool When `index` is None, this argument controls whether a new object is created or a view of an existing object is returned. index : int or None Index in the internal plots array. Attributes ---------- id : int Unique ID of the plot. pixels : tuple of int Plot image dimensions as ``(width, height)``. color_by : int Coloring mode. Use :attr:`COLOR_BY_MATERIAL` or :attr:`COLOR_BY_CELL`. camera_position : tuple of float Camera position as ``(x, y, z)``. look_at : tuple of float Point the camera is aimed at as ``(x, y, z)``. up : tuple of float Up direction as ``(x, y, z)``. light_position : tuple of float Position of the light source as ``(x, y, z)``. fov : float Horizontal field-of-view angle in degrees. diffuse_fraction : float Fraction of reflected light treated as diffuse (0 to 1). """ COLOR_BY_MATERIAL = 0 COLOR_BY_CELL = 1 __instances = WeakValueDictionary() def __new__(cls, uid=None, new=True, index=None): mapping = plots if index is None: if new: if uid is not None and uid in mapping: raise AllocationError( f'A plot with ID={uid} has already been allocated.' ) index = c_int32() _dll.openmc_solidraytrace_plot_create(index) index = index.value else: index = mapping[uid]._index if index not in cls.__instances: instance = super().__new__(cls) instance._index = index if uid is not None: instance.id = uid cls.__instances[index] = instance return cls.__instances[index] def __init__(self, uid=None, new=True, index=None): super().__init__(uid, new, index) @property def id(self): plot_id = c_int32() _dll.openmc_plot_get_id(self._index, plot_id) return plot_id.value @id.setter def id(self, plot_id): _dll.openmc_plot_set_id(self._index, plot_id) @staticmethod def _get_xyz(getter, index): x = c_double() y = c_double() z = c_double() getter(index, x, y, z) return (x.value, y.value, z.value) @staticmethod def _set_xyz(setter, index, xyz): x, y, z = xyz setter(index, float(x), float(y), float(z)) @property def pixels(self): width = c_int32() height = c_int32() _dll.openmc_solidraytrace_plot_get_pixels(self._index, width, height) return (width.value, height.value) @pixels.setter def pixels(self, pixels): width, height = pixels _dll.openmc_solidraytrace_plot_set_pixels( self._index, int(width), int(height)) @property def color_by(self): color_by = c_int32() _dll.openmc_solidraytrace_plot_get_color_by(self._index, color_by) return color_by.value @color_by.setter def color_by(self, color_by): _dll.openmc_solidraytrace_plot_set_color_by(self._index, int(color_by)) def set_default_colors(self): _dll.openmc_solidraytrace_plot_set_default_colors(self._index) def set_all_opaque(self): _dll.openmc_solidraytrace_plot_set_all_opaque(self._index) def set_visibility(self, domain_id, visible): _dll.openmc_solidraytrace_plot_set_opaque( self._index, int(domain_id), bool(visible) ) def set_color(self, domain_id, color): r, g, b = [int(c) for c in color] _dll.openmc_solidraytrace_plot_set_color( self._index, int(domain_id), r, g, b) @property def camera_position(self): return self._get_xyz(_dll.openmc_solidraytrace_plot_get_camera_position, self._index) @camera_position.setter def camera_position(self, position): self._set_xyz(_dll.openmc_solidraytrace_plot_set_camera_position, self._index, position) @property def look_at(self): return self._get_xyz(_dll.openmc_solidraytrace_plot_get_look_at, self._index) @look_at.setter def look_at(self, position): self._set_xyz(_dll.openmc_solidraytrace_plot_set_look_at, self._index, position) @property def up(self): return self._get_xyz(_dll.openmc_solidraytrace_plot_get_up, self._index) @up.setter def up(self, direction): self._set_xyz(_dll.openmc_solidraytrace_plot_set_up, self._index, direction) @property def light_position(self): return self._get_xyz(_dll.openmc_solidraytrace_plot_get_light_position, self._index) @light_position.setter def light_position(self, position): self._set_xyz(_dll.openmc_solidraytrace_plot_set_light_position, self._index, position) @property def fov(self): fov = c_double() _dll.openmc_solidraytrace_plot_get_fov(self._index, fov) return fov.value @fov.setter def fov(self, fov): _dll.openmc_solidraytrace_plot_set_fov(self._index, float(fov)) def update_view(self): _dll.openmc_solidraytrace_plot_update_view(self._index) def create_image(self): width, height = self.pixels image = np.zeros((height, width, 3), dtype=np.uint8) _dll.openmc_solidraytrace_plot_create_image( self._index, image.ctypes.data_as(POINTER(c_uint8)), width, height ) return image def get_color(self, domain_id): r = c_uint8() g = c_uint8() b = c_uint8() _dll.openmc_solidraytrace_plot_get_color( self._index, int(domain_id), r, g, b) return int(r.value), int(g.value), int(b.value) @property def diffuse_fraction(self): value = c_double() _dll.openmc_solidraytrace_plot_get_diffuse_fraction(self._index, value) return value.value @diffuse_fraction.setter def diffuse_fraction(self, value): _dll.openmc_solidraytrace_plot_set_diffuse_fraction( self._index, float(value)) # Backward-compatible setter aliases def set_pixels(self, width, height): self.pixels = (width, height) def set_color_by(self, color_by): self.color_by = color_by def set_camera_position(self, x, y, z): self.camera_position = (x, y, z) def set_look_at(self, x, y, z): self.look_at = (x, y, z) def set_up(self, x, y, z): self.up = (x, y, z) def set_light_position(self, x, y, z): self.light_position = (x, y, z) def set_fov(self, fov): self.fov = fov def set_diffuse_fraction(self, value): self.diffuse_fraction = value
class _PlotMapping(Mapping): def __getitem__(self, key): index = c_int32() try: _dll.openmc_get_plot_index(key, index) except (AllocationError, InvalidIDError) as e: raise KeyError(str(e)) return SolidRayTracePlot(index=index.value) def __iter__(self): for i in range(len(self)): yield SolidRayTracePlot(index=i).id def __len__(self): return _dll.openmc_plots_size() def __repr__(self): return repr(dict(self)) plots = _PlotMapping()