From e6a36ff7960e71cbe3a83a53138b61b61ad32bc5 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Tue, 23 Jan 2024 06:54:08 -0600 Subject: [PATCH] Add C API function for getting mesh bins for rasterized plot (#2854) --- openmc/lib/mesh.py | 48 ++++++++++++++++++++++++++++- setup.py | 2 +- src/mesh.cpp | 58 ++++++++++++++++++++++++++++++++++++ tests/unit_tests/test_lib.py | 25 ++++++++++++++++ 4 files changed, 131 insertions(+), 2 deletions(-) diff --git a/openmc/lib/mesh.py b/openmc/lib/mesh.py index 8a9953754e..599b750fb0 100644 --- a/openmc/lib/mesh.py +++ b/openmc/lib/mesh.py @@ -2,7 +2,7 @@ from collections.abc import Mapping from ctypes import (c_int, c_int32, c_char_p, c_double, POINTER, Structure, create_string_buffer, c_uint64, c_size_t) from random import getrandbits -from typing import Optional, List, Tuple +from typing import Optional, List, Tuple, Sequence from weakref import WeakValueDictionary import numpy as np @@ -13,6 +13,7 @@ from . import _dll from .core import _FortranObjectWithID from .error import _error_handler from .material import Material +from .plot import _Position __all__ = ['RegularMesh', 'RectilinearMesh', 'CylindricalMesh', 'SphericalMesh', 'UnstructuredMesh', 'meshes'] @@ -43,6 +44,11 @@ _dll.openmc_mesh_material_volumes.argtypes = [ POINTER(c_int), POINTER(c_uint64)] _dll.openmc_mesh_material_volumes.restype = c_int _dll.openmc_mesh_material_volumes.errcheck = _error_handler +_dll.openmc_mesh_get_plot_bins.argtypes = [ + c_int32, _Position, _Position, c_int, POINTER(c_int), POINTER(c_int32) +] +_dll.openmc_mesh_get_plot_bins.restype = c_int +_dll.openmc_mesh_get_plot_bins.errcheck = _error_handler _dll.openmc_get_mesh_index.argtypes = [c_int32, POINTER(c_int32)] _dll.openmc_get_mesh_index.restype = c_int _dll.openmc_get_mesh_index.errcheck = _error_handler @@ -203,6 +209,46 @@ class Mesh(_FortranObjectWithID): ]) return volumes + def get_plot_bins( + self, + origin: Sequence[float], + width: Sequence[float], + basis: str, + pixels: Sequence[int] + ) -> np.ndarray: + """Get mesh bin indices for a rasterized plot. + + .. versionadded:: 0.14.1 + + Parameters + ---------- + origin : iterable of float + Origin of the plotting view. Should have length 3. + width : iterable of float + Width of the plotting view. Should have length 2. + basis : {'xy', 'xz', 'yz'} + Plotting basis. + pixels : iterable of int + Number of pixels in each direction. Should have length 2. + + Returns + ------- + 2D numpy array with mesh bin indices corresponding to each pixel within + the plotting view. + + """ + origin = _Position(*origin) + width = _Position(*width) + basis = {'xy': 1, 'xz': 2, 'yz': 3}[basis] + pixel_array = (c_int*2)(*pixels) + img_data = np.zeros((pixels[1], pixels[0]), dtype=np.dtype('int32')) + + _dll.openmc_mesh_get_plot_bins( + self._index, origin, width, basis, pixel_array, + img_data.ctypes.data_as(POINTER(c_int32)) + ) + return img_data + class RegularMesh(Mesh): """RegularMesh stored internally. diff --git a/setup.py b/setup.py index eeb0b3f293..84a2ad367e 100755 --- a/setup.py +++ b/setup.py @@ -64,7 +64,7 @@ kwargs = { # Dependencies 'python_requires': '>=3.7', 'install_requires': [ - 'numpy>=1.9', 'h5py', 'scipy', 'ipython', 'matplotlib', + 'numpy>=1.9', 'h5py', 'scipy<1.12', 'ipython', 'matplotlib', 'pandas', 'lxml', 'uncertainties' ], 'extras_require': { diff --git a/src/mesh.cpp b/src/mesh.cpp index b70692347f..f343b27e33 100644 --- a/src/mesh.cpp +++ b/src/mesh.cpp @@ -29,6 +29,7 @@ #include "openmc/message_passing.h" #include "openmc/openmp_interface.h" #include "openmc/particle_data.h" +#include "openmc/plot.h" #include "openmc/random_dist.h" #include "openmc/search.h" #include "openmc/settings.h" @@ -1898,6 +1899,63 @@ extern "C" int openmc_mesh_material_volumes(int32_t index, int n_sample, return (n == -1) ? OPENMC_E_ALLOCATE : 0; } +extern "C" int openmc_mesh_get_plot_bins(int32_t index, Position origin, + Position width, int basis, int* pixels, int32_t* data) +{ + if (int err = check_mesh(index)) + return err; + const auto& mesh = model::meshes[index].get(); + + int pixel_width = pixels[0]; + int pixel_height = pixels[1]; + + // get pixel size + double in_pixel = (width[0]) / static_cast(pixel_width); + double out_pixel = (width[1]) / static_cast(pixel_height); + + // setup basis indices and initial position centered on pixel + int in_i, out_i; + Position xyz = origin; + enum class PlotBasis { xy = 1, xz = 2, yz = 3 }; + PlotBasis basis_enum = static_cast(basis); + switch (basis_enum) { + 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; + default: + UNREACHABLE(); + } + + // set initial position + xyz[in_i] = origin[in_i] - width[0] / 2. + in_pixel / 2.; + xyz[out_i] = origin[out_i] + width[1] / 2. - out_pixel / 2.; + +#pragma omp parallel + { + Position r = xyz; + +#pragma omp for + for (int y = 0; y < pixel_height; y++) { + r[out_i] = xyz[out_i] - out_pixel * y; + for (int x = 0; x < pixel_width; x++) { + r[in_i] = xyz[in_i] + in_pixel * x; + data[pixel_width * y + x] = mesh->get_bin(r); + } + } + } + + return 0; +} + //! Get the dimension of a regular mesh extern "C" int openmc_regular_mesh_get_dimension( int32_t index, int** dims, int* n) diff --git a/tests/unit_tests/test_lib.py b/tests/unit_tests/test_lib.py index 8aa2ea6e9c..6ae13ef76e 100644 --- a/tests/unit_tests/test_lib.py +++ b/tests/unit_tests/test_lib.py @@ -605,6 +605,31 @@ def test_regular_mesh(lib_init): assert sum(f[1] for f in elem_vols) == pytest.approx(1.26 * 1.26, 1e-2) +def test_regular_mesh_get_plot_bins(lib_init): + mesh: openmc.lib.RegularMesh = openmc.lib.meshes[2] + mesh.dimension = (2, 2, 1) + mesh.set_parameters(lower_left=(-1.0, -1.0, -0.5), + upper_right=(1.0, 1.0, 0.5)) + + # Get bins for a plot view covering only a single mesh bin + mesh_bins = mesh.get_plot_bins((-0.5, -0.5, 0.), (0.1, 0.1), 'xy', (20, 20)) + assert (mesh_bins == 0).all() + mesh_bins = mesh.get_plot_bins((0.5, 0.5, 0.), (0.1, 0.1), 'xy', (20, 20)) + assert (mesh_bins == 3).all() + + # Get bins for a plot view covering all mesh bins. Note that the y direction + # (first dimension) is flipped for plotting purposes + mesh_bins = mesh.get_plot_bins((0., 0., 0.), (2., 2.), 'xy', (20, 20)) + assert (mesh_bins[:10, :10] == 2).all() + assert (mesh_bins[:10, 10:] == 3).all() + assert (mesh_bins[10:, :10] == 0).all() + assert (mesh_bins[10:, 10:] == 1).all() + + # Get bins for a plot view outside of the mesh + mesh_bins = mesh.get_plot_bins((100., 100., 0.), (2., 2.), 'xy', (20, 20)) + assert (mesh_bins == -1).all() + + def test_rectilinear_mesh(lib_init): mesh = openmc.lib.RectilinearMesh() x_grid = [-10., 0., 10.]