From e060534ff1d79885ef2b20ae169a8b9b5e63c3c5 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Wed, 26 Feb 2025 08:14:53 -0600 Subject: [PATCH] Compute material volumes in mesh elements based on raytracing (#3129) Co-authored-by: Olek <45364492+yardasol@users.noreply.github.com> Co-authored-by: Patrick Shriwise --- docs/source/pythonapi/base.rst | 1 + include/openmc/bounding_box.h | 4 + include/openmc/capi.h | 4 +- include/openmc/mesh.h | 95 +++++-- openmc/lib/mesh.py | 131 ++++++---- openmc/mesh.py | 227 +++++++++++++--- src/mesh.cpp | 465 ++++++++++++++++++++++++++------- src/particle.cpp | 1 - tests/unit_tests/test_lib.py | 36 +-- tests/unit_tests/test_mesh.py | 88 ++++++- 10 files changed, 814 insertions(+), 238 deletions(-) diff --git a/docs/source/pythonapi/base.rst b/docs/source/pythonapi/base.rst index bfd942727..e5b5b17c3 100644 --- a/docs/source/pythonapi/base.rst +++ b/docs/source/pythonapi/base.rst @@ -152,6 +152,7 @@ Constructing Tallies openmc.CylindricalMesh openmc.SphericalMesh openmc.UnstructuredMesh + openmc.MeshMaterialVolumes openmc.Trigger openmc.TallyDerivative openmc.Tally diff --git a/include/openmc/bounding_box.h b/include/openmc/bounding_box.h index 40f603583..d02d92cb4 100644 --- a/include/openmc/bounding_box.h +++ b/include/openmc/bounding_box.h @@ -4,6 +4,7 @@ #include // for min, max #include "openmc/constants.h" +#include "openmc/position.h" namespace openmc { @@ -54,6 +55,9 @@ struct BoundingBox { zmax = std::max(zmax, other.zmax); return *this; } + + inline Position min() const { return {xmin, ymin, zmin}; } + inline Position max() const { return {xmax, ymax, zmax}; } }; } // namespace openmc diff --git a/include/openmc/capi.h b/include/openmc/capi.h index 8edd99c07..3802692c3 100644 --- a/include/openmc/capi.h +++ b/include/openmc/capi.h @@ -110,8 +110,8 @@ int openmc_mesh_get_id(int32_t index, int32_t* id); int openmc_mesh_set_id(int32_t index, int32_t id); int openmc_mesh_get_n_elements(int32_t index, size_t* n); int openmc_mesh_get_volumes(int32_t index, double* volumes); -int openmc_mesh_material_volumes(int32_t index, int n_sample, int bin, - int result_size, void* result, int* hits, uint64_t* seed); +int openmc_mesh_material_volumes(int32_t index, int nx, int ny, int nz, + int max_mats, int32_t* materials, double* volumes); int openmc_meshsurface_filter_get_mesh(int32_t index, int32_t* index_mesh); int openmc_meshsurface_filter_set_mesh(int32_t index, int32_t index_mesh); int openmc_new_filter(const char* type, int32_t* index); diff --git a/include/openmc/mesh.h b/include/openmc/mesh.h index 3aef49c7e..943e49a5f 100644 --- a/include/openmc/mesh.h +++ b/include/openmc/mesh.h @@ -69,14 +69,70 @@ extern const libMesh::Parallel::Communicator* libmesh_comm; } // namespace settings #endif +//============================================================================== +//! Helper class for keeping track of volume for each material in a mesh element +// +//! This class is used in Mesh::material_volumes to manage for each mesh element +//! a list of (material, volume) pairs. The openmc.lib.Mesh class allocates two +//! 2D arrays, one for materials and one for volumes. Because we don't know a +//! priori how many materials there are in each element but at the same time we +//! can't dynamically size an array at runtime for performance reasons, we +//! assume a maximum number of materials per element. For each element, the set +//! of material indices are stored in a hash table with twice as many slots as +//! the assumed maximum number of materials per element. Collision resolution is +//! handled by open addressing with linear probing. +//============================================================================== + +namespace detail { + +class MaterialVolumes { +public: + MaterialVolumes(int32_t* mats, double* vols, int table_size) + : materials_(mats), volumes_(vols), table_size_(table_size) + {} + + //! Add volume for a given material in a mesh element + // + //! \param[in] index_elem Index of the mesh element + //! \param[in] index_material Index of the material within the model + //! \param[in] volume Volume to add + void add_volume(int index_elem, int index_material, double volume); + void add_volume_unsafe(int index_elem, int index_material, double volume); + + // Accessors + int32_t& materials(int i, int j) { return materials_[i * table_size_ + j]; } + const int32_t& materials(int i, int j) const + { + return materials_[i * table_size_ + j]; + } + + double& volumes(int i, int j) { return volumes_[i * table_size_ + j]; } + const double& volumes(int i, int j) const + { + return volumes_[i * table_size_ + j]; + } + + bool table_full() const { return table_full_; } + +private: + int32_t* materials_; //!< material index (bins, table_size) + double* volumes_; //!< volume in [cm^3] (bins, table_size) + int table_size_; //!< Size of hash table for each mesh element + bool table_full_ {false}; //!< Whether the hash table is full + + // Value used to indicate an empty slot in the hash table. We use -2 because + // the value -1 is used to indicate a void material. + static constexpr int EMPTY {-2}; +}; + +} // namespace detail + +//============================================================================== +//! Base mesh class +//============================================================================== + class Mesh { public: - // Types, aliases - struct MaterialVolume { - int32_t material; //!< material index - double volume; //!< volume in [cm^3] - }; - // Constructors and destructor Mesh() = default; Mesh(pugi::xml_node node); @@ -172,24 +228,17 @@ public: virtual std::string get_mesh_type() const = 0; - //! Determine volume of materials within a single mesh elemenet + //! Determine volume of materials within each mesh element // - //! \param[in] n_sample Number of samples within each element - //! \param[in] bin Index of mesh element - //! \param[out] Array of (material index, volume) for desired element - //! \param[inout] seed Pseudorandom number seed - //! \return Number of materials within element - int material_volumes( - int n_sample, int bin, span volumes, uint64_t* seed) const; - - //! Determine volume of materials within a single mesh elemenet - // - //! \param[in] n_sample Number of samples within each element - //! \param[in] bin Index of mesh element - //! \param[inout] seed Pseudorandom number seed - //! \return Vector of (material index, volume) for desired element - vector material_volumes( - int n_sample, int bin, uint64_t* seed) const; + //! \param[in] nx Number of samples in x direction + //! \param[in] ny Number of samples in y direction + //! \param[in] nz Number of samples in z direction + //! \param[in] max_materials Maximum number of materials in a single mesh + //! element + //! \param[inout] materials Array storing material indices + //! \param[inout] volumes Array storing volumes + void material_volumes(int nx, int ny, int nz, int max_materials, + int32_t* materials, double* volumes) const; //! Determine bounding box of mesh // diff --git a/openmc/lib/mesh.py b/openmc/lib/mesh.py index 16bec0198..3a720f98a 100644 --- a/openmc/lib/mesh.py +++ b/openmc/lib/mesh.py @@ -1,7 +1,7 @@ from collections.abc import Mapping, Sequence -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 ctypes import (c_int, c_int32, c_char_p, c_double, POINTER, + create_string_buffer, c_size_t) +from math import sqrt import sys from weakref import WeakValueDictionary @@ -10,24 +10,20 @@ from numpy.ctypeslib import as_array from ..exceptions import AllocationError, InvalidIDError from . import _dll -from .core import _FortranObjectWithID +from .core import _FortranObjectWithID, quiet_dll from .error import _error_handler -from .material import Material from .plot import _Position from ..bounding_box import BoundingBox +from ..mesh import MeshMaterialVolumes __all__ = [ 'Mesh', 'RegularMesh', 'RectilinearMesh', 'CylindricalMesh', - 'SphericalMesh', 'UnstructuredMesh', 'meshes' + 'SphericalMesh', 'UnstructuredMesh', 'meshes', 'MeshMaterialVolumes' ] -class _MaterialVolume(Structure): - _fields_ = [ - ("material", c_int32), - ("volume", c_double) - ] - +arr_2d_int32 = np.ctypeslib.ndpointer(dtype=np.int32, ndim=2, flags='CONTIGUOUS') +arr_2d_double = np.ctypeslib.ndpointer(dtype=np.double, ndim=2, flags='CONTIGUOUS') # Mesh functions _dll.openmc_extend_meshes.argtypes = [c_int32, c_char_p, POINTER(c_int32), @@ -51,8 +47,7 @@ _dll.openmc_mesh_bounding_box.argtypes = [ _dll.openmc_mesh_bounding_box.restype = c_int _dll.openmc_mesh_bounding_box.errcheck = _error_handler _dll.openmc_mesh_material_volumes.argtypes = [ - c_int32, c_int, c_int, c_int, POINTER(_MaterialVolume), - POINTER(c_int), POINTER(c_uint64)] + c_int32, c_int, c_int, c_int, c_int, arr_2d_int32, arr_2d_double] _dll.openmc_mesh_material_volumes.restype = c_int _dll.openmc_mesh_material_volumes.errcheck = _error_handler _dll.openmc_mesh_get_plot_bins.argtypes = [ @@ -190,58 +185,81 @@ class Mesh(_FortranObjectWithID): def material_volumes( self, - n_samples: int = 10_000, - prn_seed: int | None = None - ) -> list[list[tuple[Material, float]]]: - """Determine volume of materials in each mesh element + n_samples: int | tuple[int, int, int] = 10_000, + max_materials: int = 4, + output: bool = True, + ) -> MeshMaterialVolumes: + """Determine volume of materials in each mesh element. + + This method works by raytracing repeatedly through the mesh to count the + estimated volume of each material in all mesh elements. Three sets of + rays are used: one set parallel to the x-axis, one parallel to the + y-axis, and one parallel to the z-axis. .. versionadded:: 0.15.0 + .. versionchanged:: 0.15.1 + Material volumes are now determined by raytracing rather than by + point sampling. + Parameters ---------- - n_samples : int - Number of samples in each mesh element - prn_seed : int - Pseudorandom number generator (PRNG) seed; if None, one will be - generated randomly. + n_samples : int or 3-tuple of int + Total number of rays to sample. The number of rays in each direction + is determined by the aspect ratio of the mesh bounding box. When + specified as a 3-tuple, it is interpreted as the number of rays in + the x, y, and z dimensions. + max_materials : int, optional + Estimated maximum number of materials in any given mesh element. + output : bool, optional + Whether or not to show output. Returns ------- - List of tuple of (material, volume) for each mesh element. Void volume - is represented by having a value of None in the first element of a - tuple. + MeshMaterialVolumes + Dictionary-like object that maps material IDs to an array of volumes + equal in size to the number of mesh elements. """ - if n_samples <= 0: - raise ValueError("Number of samples must be positive") - if prn_seed is None: - prn_seed = getrandbits(63) - prn_seed = c_uint64(prn_seed) + if isinstance(n_samples, int): + # Determine number of rays in each direction based on aspect ratios + # and using the relation (nx*ny + ny*nz + nx*nz) = n_samples + width_x, width_y, width_z = self.bounding_box.width + ax = width_x / width_z + ay = width_y / width_z + f = sqrt(n_samples/(ax*ay + ax + ay)) + nx = round(f * ax) + ny = round(f * ay) + nz = round(f) + else: + nx, ny, nz = n_samples - # Preallocate space for MaterialVolume results - size = 16 - result = (_MaterialVolume * size)() + # Value indicating an empty slot in the hash table (matches C++) + EMPTY_SLOT = -2 - hits = c_int() # Number of materials hit in a given element - volumes = [] - for i_element in range(self.n_elements): - while True: - try: + # Preallocate arrays for material indices and volumes + n = self.n_elements + slot_factor = 2 + table_size = slot_factor*max_materials + materials = np.full((n, table_size), EMPTY_SLOT, dtype=np.int32) + volumes = np.zeros((n, table_size), dtype=np.float64) + + # Run material volume calculation + while True: + try: + with quiet_dll(output): _dll.openmc_mesh_material_volumes( - self._index, n_samples, i_element, size, result, hits, prn_seed) - except AllocationError: - # Increase size of result array and try again - size *= 2 - result = (_MaterialVolume * size)() - else: - # If no error, break out of loop - break + self._index, nx, ny, nz, table_size, materials, volumes) + except AllocationError: + # Increase size of result array and try again + table_size *= 2 + materials = np.full((n, table_size), EMPTY_SLOT, dtype=np.int32) + volumes = np.zeros((n, table_size), dtype=np.float64) + else: + # If no error, break out of loop + break - volumes.append([ - (Material(index=r.material), r.volume) - for r in result[:hits.value] - ]) - return volumes + return MeshMaterialVolumes(materials, volumes) def get_plot_bins( self, @@ -306,7 +324,7 @@ class RegularMesh(Mesh): The lower-left corner of the structured mesh. If only two coordinate are given, it is assumed that the mesh is an x-y mesh. upper_right : numpy.ndarray - The upper-right corner of the structrued mesh. If only two coordinate + The upper-right corner of the structured mesh. If only two coordinate are given, it is assumed that the mesh is an x-y mesh. width : numpy.ndarray The width of mesh cells in each direction. @@ -395,7 +413,7 @@ class RectilinearMesh(Mesh): lower_left : numpy.ndarray The lower-left corner of the structured mesh. upper_right : numpy.ndarray - The upper-right corner of the structrued mesh. + The upper-right corner of the structured mesh. width : numpy.ndarray The width of mesh cells in each direction. n_elements : int @@ -500,7 +518,7 @@ class CylindricalMesh(Mesh): lower_left : numpy.ndarray The lower-left corner of the structured mesh. upper_right : numpy.ndarray - The upper-right corner of the structrued mesh. + The upper-right corner of the structured mesh. width : numpy.ndarray The width of mesh cells in each direction. n_elements : int @@ -605,7 +623,7 @@ class SphericalMesh(Mesh): lower_left : numpy.ndarray The lower-left corner of the structured mesh. upper_right : numpy.ndarray - The upper-right corner of the structrued mesh. + The upper-right corner of the structured mesh. width : numpy.ndarray The width of mesh cells in each direction. n_elements : int @@ -719,7 +737,6 @@ class _MeshMapping(Mapping): raise KeyError(str(e)) return _get_mesh(index.value) - def __iter__(self): for i in range(len(self)): yield _get_mesh(i).id diff --git a/openmc/mesh.py b/openmc/mesh.py index 4c594c9c6..08025f374 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -1,7 +1,7 @@ from __future__ import annotations import warnings from abc import ABC, abstractmethod -from collections.abc import Iterable, Sequence +from collections.abc import Iterable, Sequence, Mapping from functools import wraps from math import pi, sqrt, atan2 from numbers import Integral, Real @@ -21,6 +21,120 @@ from .surface import _BOUNDARY_TYPES from .utility_funcs import input_path +class MeshMaterialVolumes(Mapping): + """Results from a material volume in mesh calculation. + + This class provides multiple ways of accessing information about material + volumes in individual mesh elements. First, the class behaves like a + dictionary that maps material IDs to an array of volumes equal in size to + the number of mesh elements. Second, the class provides a :meth:`by_element` + method that gives all the material volumes for a specific mesh element. + + .. versionadded:: 0.15.1 + + Parameters + ---------- + materials : numpy.ndarray + Array of shape (elements, max_materials) storing material IDs + volumes : numpy.ndarray + Array of shape (elements, max_materials) storing material volumes + + See Also + -------- + openmc.MeshBase.material_volumes + + Examples + -------- + If you want to get the volume of a specific material in every mesh element, + index the object with the material ID: + + >>> volumes = mesh.material_volumes(...) + >>> volumes + {1: <32121 nonzero volumes> + 2: <338186 nonzero volumes> + 3: <49120 nonzero volumes>} + + If you want the volume of all materials in a specific mesh element, use the + :meth:`by_element` method: + + >>> volumes = mesh.material_volumes(...) + >>> volumes.by_element(42) + [(2, 31.87963824195591), (1, 6.129949130817542)] + + """ + def __init__(self, materials: np.ndarray, volumes: np.ndarray): + self._materials = materials + self._volumes = volumes + + @property + def num_elements(self) -> int: + return self._volumes.shape[0] + + def __iter__(self): + for mat in np.unique(self._materials): + if mat > 0: + yield mat + + def __len__(self) -> int: + return (np.unique(self._materials) > 0).sum() + + def __repr__(self) -> str: + ids, counts = np.unique(self._materials, return_counts=True) + return '{' + '\n '.join( + f'{id}: <{count} nonzero volumes>' for id, count in zip(ids, counts) if id > 0) + '}' + + def __getitem__(self, material_id: int) -> np.ndarray: + volumes = np.zeros(self.num_elements) + for i in range(self._volumes.shape[1]): + indices = (self._materials[:, i] == material_id) + volumes[indices] = self._volumes[indices, i] + return volumes + + def by_element(self, index_elem: int) -> list[tuple[int | None, float]]: + """Get a list of volumes for each material within a specific element. + + Parameters + ---------- + index_elem : int + Mesh element index + + Returns + ------- + list of tuple of (material ID, volume) + + """ + table_size = self._volumes.shape[1] + return [ + (m if m > -1 else None, self._volumes[index_elem, i]) + for i in range(table_size) + if (m := self._materials[index_elem, i]) != -2 + ] + + def save(self, filename: PathLike): + """Save material volumes to a .npz file. + + Parameters + ---------- + filename : path-like + Filename where data will be saved + """ + np.savez_compressed( + filename, materials=self._materials, volumes=self._volumes) + + @classmethod + def from_npz(cls, filename: PathLike) -> MeshMaterialVolumes: + """Generate material volumes from a .npz file + + Parameters + ---------- + filename : path-like + File where data will be read from + + """ + filedata = np.load(filename) + return cls(filedata['materials'], filedata['volumes']) + + class MeshBase(IDManagerMixin, ABC): """A mesh that partitions geometry for tallying purposes. @@ -170,8 +284,7 @@ class MeshBase(IDManagerMixin, ABC): def get_homogenized_materials( self, model: openmc.Model, - n_samples: int = 10_000, - prn_seed: int | None = None, + n_samples: int | tuple[int, int, int] = 10_000, include_void: bool = True, **kwargs ) -> list[openmc.Material]: @@ -184,15 +297,15 @@ class MeshBase(IDManagerMixin, ABC): model : openmc.Model Model containing materials to be homogenized and the associated geometry. - n_samples : int - Number of samples in each mesh element. - prn_seed : int, optional - Pseudorandom number generator (PRNG) seed; if None, one will be - generated randomly. + n_samples : int or 2-tuple of int + Total number of rays to sample. The number of rays in each direction + is determined by the aspect ratio of the mesh bounding box. When + specified as a 3-tuple, it is interpreted as the number of rays in + the x, y, and z dimensions. include_void : bool, optional Whether homogenization should include voids. **kwargs - Keyword-arguments passed to :func:`openmc.lib.init`. + Keyword-arguments passed to :meth:`MeshBase.material_volumes`. Returns ------- @@ -200,34 +313,8 @@ class MeshBase(IDManagerMixin, ABC): Homogenized material in each mesh element """ - import openmc.lib - - with change_directory(tmpdir=True): - # In order to get mesh into model, we temporarily replace the - # tallies with a single mesh tally using the current mesh - original_tallies = model.tallies - new_tally = openmc.Tally() - new_tally.filters = [openmc.MeshFilter(self)] - new_tally.scores = ['flux'] - model.tallies = [new_tally] - - # Export model to XML - model.export_to_model_xml() - - # Get material volume fractions - openmc.lib.init(**kwargs) - mesh = openmc.lib.tallies[new_tally.id].filters[0].mesh - mat_volume_by_element = [ - [ - (mat.id if mat is not None else None, volume) - for mat, volume in mat_volume_list - ] - for mat_volume_list in mesh.material_volumes(n_samples, prn_seed) - ] - openmc.lib.finalize() - - # Restore original tallies - model.tallies = original_tallies + vols = self.material_volumes(model, n_samples, **kwargs) + mat_volume_by_element = [vols.by_element(i) for i in range(vols.num_elements)] # Create homogenized material for each element materials = model.geometry.get_all_materials() @@ -274,6 +361,72 @@ class MeshBase(IDManagerMixin, ABC): return homogenized_materials + def material_volumes( + self, + model: openmc.Model, + n_samples: int | tuple[int, int, int] = 10_000, + max_materials: int = 4, + **kwargs + ) -> MeshMaterialVolumes: + """Determine volume of materials in each mesh element. + + This method works by raytracing repeatedly through the mesh to count the + estimated volume of each material in all mesh elements. Three sets of + rays are used: one set parallel to the x-axis, one parallel to the + y-axis, and one parallel to the z-axis. + + .. versionadded:: 0.15.1 + + Parameters + ---------- + model : openmc.Model + Model containing materials. + n_samples : int or 3-tuple of int + Total number of rays to sample. The number of rays in each direction + is determined by the aspect ratio of the mesh bounding box. When + specified as a 3-tuple, it is interpreted as the number of rays in + the x, y, and z dimensions. + max_materials : int, optional + Estimated maximum number of materials in any given mesh element. + **kwargs : dict + Keyword arguments passed to :func:`openmc.lib.init` + + Returns + ------- + Dictionary-like object that maps material IDs to an array of volumes + equal in size to the number of mesh elements. + + """ + import openmc.lib + + with change_directory(tmpdir=True): + # In order to get mesh into model, we temporarily replace the + # tallies with a single mesh tally using the current mesh + original_tallies = model.tallies + new_tally = openmc.Tally() + new_tally.filters = [openmc.MeshFilter(self)] + new_tally.scores = ['flux'] + model.tallies = [new_tally] + + # Export model to XML + model.export_to_model_xml() + + # Get material volume fractions + kwargs.setdefault('output', True) + if 'args' in kwargs: + kwargs['args'] = ['-c'] + kwargs['args'] + kwargs.setdefault('args', ['-c']) + openmc.lib.init(**kwargs) + mesh = openmc.lib.tallies[new_tally.id].filters[0].mesh + volumes = mesh.material_volumes( + n_samples, max_materials, output=kwargs['output']) + openmc.lib.finalize() + + # Restore original tallies + model.tallies = original_tallies + + return volumes + class StructuredMesh(MeshBase): """A base class for structured mesh functionality diff --git a/src/mesh.cpp b/src/mesh.cpp index 3e1f25c4b..a24b51a0b 100644 --- a/src/mesh.cpp +++ b/src/mesh.cpp @@ -6,10 +6,15 @@ #include // for size_t #include +#ifdef _MSC_VER +#include // for _InterlockedCompareExchange +#endif + #ifdef OPENMC_MPI #include "mpi.h" #endif +#include "xtensor/xadapt.hpp" #include "xtensor/xbuilder.hpp" #include "xtensor/xeval.hpp" #include "xtensor/xmath.hpp" @@ -29,13 +34,16 @@ #include "openmc/memory.h" #include "openmc/message_passing.h" #include "openmc/openmp_interface.h" +#include "openmc/output.h" #include "openmc/particle_data.h" #include "openmc/plot.h" #include "openmc/random_dist.h" #include "openmc/search.h" #include "openmc/settings.h" +#include "openmc/string_utils.h" #include "openmc/tallies/filter.h" #include "openmc/tallies/tally.h" +#include "openmc/timer.h" #include "openmc/volume_calc.h" #include "openmc/xml_interface.h" @@ -102,6 +110,113 @@ inline bool check_intersection_point(double x1, double x0, double y1, double y0, return false; } +//! Atomic compare-and-swap for signed 32-bit integer +// +//! \param[in,out] ptr Pointer to value to update +//! \param[in] expected Value to compare to +//! \param[in] desired If comparison is successful, value to update to +//! \return True if the comparison was successful and the value was updated +inline bool atomic_cas_int32(int32_t* ptr, int32_t& expected, int32_t desired) +{ +#if defined(__GNUC__) || defined(__clang__) + // For gcc/clang, use the __atomic_compare_exchange_n intrinsic + return __atomic_compare_exchange_n( + ptr, &expected, desired, false, __ATOMIC_SEQ_CST, __ATOMIC_SEQ_CST); + +#elif defined(_MSC_VER) + // For MSVC, use the _InterlockedCompareExchange intrinsic + int32_t old_val = + _InterlockedCompareExchange(reinterpret_cast(ptr), + static_cast(desired), static_cast(expected)); + return (old_val == expected); + +#else +#error "No compare-and-swap implementation available for this compiler." +#endif +} + +namespace detail { + +//============================================================================== +// MaterialVolumes implementation +//============================================================================== + +void MaterialVolumes::add_volume( + int index_elem, int index_material, double volume) +{ + // This method handles adding elements to the materials hash table, + // implementing open addressing with linear probing. Consistency across + // multiple threads is handled by with an atomic compare-and-swap operation. + // Ideally, we would use #pragma omp atomic compare, but it was introduced in + // OpenMP 5.1 and is not widely supported yet. + + // Loop for linear probing + for (int attempt = 0; attempt < table_size_; ++attempt) { + // Determine slot to check + int slot = (index_material + attempt) % table_size_; + int32_t* slot_ptr = &this->materials(index_elem, slot); + + // Non-atomic read of current material + int32_t current_val = *slot_ptr; + + // Found the desired material; accumulate volume + if (current_val == index_material) { +#pragma omp atomic + this->volumes(index_elem, slot) += volume; + return; + } + + // Slot appears to be empty; attempt to claim + if (current_val == EMPTY) { + // Attempt compare-and-swap from EMPTY to index_material + int32_t expected_val = EMPTY; + bool claimed_slot = + atomic_cas_int32(slot_ptr, expected_val, index_material); + + // If we claimed the slot or another thread claimed it but the same + // material was inserted, proceed to accumulate + if (claimed_slot || (expected_val == index_material)) { +#pragma omp atomic + this->volumes(index_elem, slot) += volume; + return; + } + } + } + + // If table is full, set a flag that can be checked later + table_full_ = true; +} + +void MaterialVolumes::add_volume_unsafe( + int index_elem, int index_material, double volume) +{ + // Linear probe + for (int attempt = 0; attempt < table_size_; ++attempt) { + int slot = (index_material + attempt) % table_size_; + + // Read current material + int32_t current_val = this->materials(index_elem, slot); + + // Found the desired material; accumulate volume + if (current_val == index_material) { + this->volumes(index_elem, slot) += volume; + return; + } + + // Claim empty slot + if (current_val == EMPTY) { + this->materials(index_elem, slot) = index_material; + this->volumes(index_elem, slot) += volume; + return; + } + } + + // If table is full, set a flag that can be checked later + table_full_ = true; +} + +} // namespace detail + //============================================================================== // Mesh implementation //============================================================================== @@ -153,90 +268,235 @@ vector Mesh::volumes() const return volumes; } -int Mesh::material_volumes( - int n_sample, int bin, span result, uint64_t* seed) const +void Mesh::material_volumes(int nx, int ny, int nz, int table_size, + int32_t* materials, double* volumes) const { - vector materials; - vector hits; + if (mpi::master) { + header("MESH MATERIAL VOLUMES CALCULATION", 7); + } + write_message(7, "Number of rays (x) = {}", nx); + write_message(7, "Number of rays (y) = {}", ny); + write_message(7, "Number of rays (z) = {}", nz); + int64_t n_total = nx * ny + ny * nz + nx * nz; + write_message(7, "Total number of rays = {}", n_total); + write_message( + 7, "Maximum number of materials per mesh element = {}", table_size); + + Timer timer; + timer.start(); + + // Create object for keeping track of materials/volumes + detail::MaterialVolumes result(materials, volumes, table_size); + + // Determine bounding box + auto bbox = this->bounding_box(); + + std::array n_rays = {nx, ny, nz}; + + // Determine effective width of rays + Position width((bbox.xmax - bbox.xmin) / nx, (bbox.ymax - bbox.ymin) / ny, + (bbox.zmax - bbox.zmin) / nz); + + // Set flag for mesh being contained within model + bool out_of_model = false; #pragma omp parallel { - vector local_materials; - vector local_hits; - GeometryState geom; + // Preallocate vector for mesh indices and length fractions and particle + std::vector bins; + std::vector length_fractions; + Particle p; -#pragma omp for - for (int i = 0; i < n_sample; ++i) { - // Get seed for i-th sample - uint64_t seed_i = future_seed(3 * i, *seed); + SourceSite site; + site.E = 1.0; + site.particle = ParticleType::neutron; - // Sample position and set geometry state - geom.r() = this->sample_element(bin, &seed_i); - geom.u() = {1., 0., 0.}; - geom.n_coord() = 1; + for (int axis = 0; axis < 3; ++axis) { + // Set starting position and direction + site.r = {0.0, 0.0, 0.0}; + site.r[axis] = bbox.min()[axis]; + site.u = {0.0, 0.0, 0.0}; + site.u[axis] = 1.0; - // If this location is not in the geometry at all, move on to next block - if (!exhaustive_find_cell(geom)) - continue; + // Determine width of rays and number of rays in other directions + int ax1 = (axis + 1) % 3; + int ax2 = (axis + 2) % 3; + double min1 = bbox.min()[ax1]; + double min2 = bbox.min()[ax2]; + double d1 = width[ax1]; + double d2 = width[ax2]; + int n1 = n_rays[ax1]; + int n2 = n_rays[ax2]; - int i_material = geom.material(); + // Divide rays in first direction over MPI processes by computing starting + // and ending indices + int min_work = n1 / mpi::n_procs; + int remainder = n1 % mpi::n_procs; + int n1_local = (mpi::rank < remainder) ? min_work + 1 : min_work; + int i1_start = mpi::rank * min_work + std::min(mpi::rank, remainder); + int i1_end = i1_start + n1_local; - // Check if this material was previously hit and if so, increment count - auto it = - std::find(local_materials.begin(), local_materials.end(), i_material); - if (it == local_materials.end()) { - local_materials.push_back(i_material); - local_hits.push_back(1); - } else { - local_hits[it - local_materials.begin()]++; + // Loop over rays on face of bounding box +#pragma omp for collapse(2) + for (int i1 = i1_start; i1 < i1_end; ++i1) { + for (int i2 = 0; i2 < n2; ++i2) { + site.r[ax1] = min1 + (i1 + 0.5) * d1; + site.r[ax2] = min2 + (i2 + 0.5) * d2; + + p.from_source(&site); + + // Determine particle's location + if (!exhaustive_find_cell(p)) { + out_of_model = true; + continue; + } + + // Set birth cell attribute + if (p.cell_born() == C_NONE) + p.cell_born() = p.lowest_coord().cell; + + // Initialize last cells from current cell + for (int j = 0; j < p.n_coord(); ++j) { + p.cell_last(j) = p.coord(j).cell; + } + p.n_coord_last() = p.n_coord(); + + while (true) { + // Ray trace from r_start to r_end + Position r0 = p.r(); + double max_distance = bbox.max()[axis] - r0[axis]; + + // Find the distance to the nearest boundary + BoundaryInfo boundary = distance_to_boundary(p); + + // Advance particle forward + double distance = std::min(boundary.distance, max_distance); + p.move_distance(distance); + + // Determine what mesh elements were crossed by particle + bins.clear(); + length_fractions.clear(); + this->bins_crossed(r0, p.r(), p.u(), bins, length_fractions); + + // Add volumes to any mesh elements that were crossed + int i_material = p.material(); + if (i_material != C_NONE) { + i_material = model::materials[i_material]->id(); + } + for (int i_bin = 0; i_bin < bins.size(); i_bin++) { + int mesh_index = bins[i_bin]; + double length = distance * length_fractions[i_bin]; + + // Add volume to result + result.add_volume(mesh_index, i_material, length * d1 * d2); + } + + if (distance == max_distance) + break; + + // cross next geometric surface + for (int j = 0; j < p.n_coord(); ++j) { + p.cell_last(j) = p.coord(j).cell; + } + p.n_coord_last() = p.n_coord(); + + // Set surface that particle is on and adjust coordinate levels + p.surface() = boundary.surface; + p.n_coord() = boundary.coord_level; + + if (boundary.lattice_translation[0] != 0 || + boundary.lattice_translation[1] != 0 || + boundary.lattice_translation[2] != 0) { + // Particle crosses lattice boundary + cross_lattice(p, boundary); + } else { + // Particle crosses surface + const auto& surf {model::surfaces[p.surface_index()].get()}; + p.cross_surface(*surf); + } + } + } } - } // omp for - - // Reduce index/hits lists from each thread into a single copy - reduce_indices_hits(local_materials, local_hits, materials, hits); - } // omp parallel - - // Advance RNG seed - advance_prn_seed(3 * n_sample, seed); - - // Make sure span passed in is large enough - if (hits.size() > result.size()) { - return -1; - } - - // Convert hits to fractions - for (int i_mat = 0; i_mat < hits.size(); ++i_mat) { - double fraction = double(hits[i_mat]) / n_sample; - result[i_mat].material = materials[i_mat]; - result[i_mat].volume = fraction * this->volume(bin); - } - return hits.size(); -} - -vector Mesh::material_volumes( - int n_sample, int bin, uint64_t* seed) const -{ - // Create result vector with space for 8 pairs - vector result; - result.reserve(8); - - int size = -1; - while (true) { - // Get material volumes - size = this->material_volumes( - n_sample, bin, {result.data(), result.data() + result.capacity()}, seed); - - // If capacity was sufficient, resize the vector and return - if (size >= 0) { - result.resize(size); - break; } - - // Otherwise, increase capacity of the vector - result.reserve(2 * result.capacity()); } - return result; + // Check for errors + if (out_of_model) { + throw std::runtime_error("Mesh not fully contained in geometry."); + } else if (result.table_full()) { + throw std::runtime_error("Maximum number of materials for mesh material " + "volume calculation insufficient."); + } + + // Compute time for raytracing + double t_raytrace = timer.elapsed(); + +#ifdef OPENMC_MPI + // Combine results from multiple MPI processes + if (mpi::n_procs > 1) { + int total = this->n_bins() * table_size; + if (mpi::master) { + // Allocate temporary buffer for receiving data + std::vector mats(total); + std::vector vols(total); + + for (int i = 1; i < mpi::n_procs; ++i) { + // Receive material indices and volumes from process i + MPI_Recv( + mats.data(), total, MPI_INT, i, i, mpi::intracomm, MPI_STATUS_IGNORE); + MPI_Recv(vols.data(), total, MPI_DOUBLE, i, i, mpi::intracomm, + MPI_STATUS_IGNORE); + + // Combine with existing results; we can call thread unsafe version of + // add_volume because each thread is operating on a different element +#pragma omp for + for (int index_elem = 0; index_elem < n_bins(); ++index_elem) { + for (int k = 0; k < table_size; ++k) { + int index = index_elem * table_size + k; + result.add_volume_unsafe(index_elem, mats[index], vols[index]); + } + } + } + } else { + // Send material indices and volumes to process 0 + MPI_Send(materials, total, MPI_INT, 0, mpi::rank, mpi::intracomm); + MPI_Send(volumes, total, MPI_DOUBLE, 0, mpi::rank, mpi::intracomm); + } + } + + // Report time for MPI communication + double t_mpi = timer.elapsed() - t_raytrace; +#else + double t_mpi = 0.0; +#endif + + // Normalize based on known volumes of elements + for (int i = 0; i < this->n_bins(); ++i) { + // Estimated total volume in element i + double volume = 0.0; + for (int j = 0; j < table_size; ++j) { + volume += result.volumes(i, j); + } + // Renormalize volumes based on known volume of element i + double norm = this->volume(i) / volume; + for (int j = 0; j < table_size; ++j) { + result.volumes(i, j) *= norm; + } + } + + // Show elapsed time + timer.stop(); + double t_total = timer.elapsed(); + double t_normalize = t_total - t_raytrace - t_mpi; + if (mpi::master) { + header("Timing Statistics", 7); + show_time("Total time elapsed", t_total); + show_time("Ray tracing", t_raytrace, 1); + show_time("Ray tracing (per ray)", t_raytrace / n_total, 1); + show_time("MPI communication", t_mpi, 1); + show_time("Normalization", t_normalize, 1); + std::fflush(stdout); + } } void Mesh::to_hdf5(hid_t group) const @@ -616,8 +876,8 @@ xt::xtensor StructuredMesh::count_sites( } // raytrace through the mesh. The template class T will do the tallying. -// A modern optimizing compiler can recognize the noop method of T and eleminate -// that call entirely. +// A modern optimizing compiler can recognize the noop method of T and +// eliminate that call entirely. template void StructuredMesh::raytrace_mesh( Position r0, Position r1, const Direction& u, T tally) const @@ -685,7 +945,8 @@ void StructuredMesh::raytrace_mesh( if (traveled_distance >= total_distance) return; - // If we have not reached r1, we have hit a surface. Tally outward current + // If we have not reached r1, we have hit a surface. Tally outward + // current tally.surface(ijk, k, distances[k].max_surface, false); // Update cell and calculate distance to next surface in k-direction. @@ -697,15 +958,16 @@ void StructuredMesh::raytrace_mesh( // Check if we have left the interior of the mesh in_mesh = ((ijk[k] >= 1) && (ijk[k] <= shape_[k])); - // If we are still inside the mesh, tally inward current for the next cell + // If we are still inside the mesh, tally inward current for the next + // cell if (in_mesh) tally.surface(ijk, k, !distances[k].max_surface, true); } else { // not inside mesh - // For all directions outside the mesh, find the distance that we need to - // travel to reach the next surface. Use the largest distance, as only - // this will cross all outer surfaces. + // For all directions outside the mesh, find the distance that we need + // to travel to reach the next surface. Use the largest distance, as + // only this will cross all outer surfaces. int k_max {0}; for (int k = 0; k < n; ++k) { if ((ijk[k] < 1 || ijk[k] > shape_[k]) && @@ -719,7 +981,8 @@ void StructuredMesh::raytrace_mesh( if (traveled_distance >= total_distance) return; - // Calculate the new cell index and update all distances to next surfaces. + // Calculate the new cell index and update all distances to next + // surfaces. ijk = get_indices(r0 + (traveled_distance + TINY_BIT) * u, in_mesh); for (int k = 0; k < n; ++k) { distances[k] = @@ -1581,7 +1844,8 @@ double SphericalMesh::find_theta_crossing( const double b = r.dot(u) * cos_t_2 - r.z * u.z; const double c = r.dot(r) * cos_t_2 - r.z * r.z; - // if factor of s^2 is zero, direction of flight is parallel to theta surface + // if factor of s^2 is zero, direction of flight is parallel to theta + // surface if (std::abs(a) < FP_PRECISION) { // if b vanishes, direction of flight is within theta surface and crossing // is not possible @@ -1589,7 +1853,8 @@ double SphericalMesh::find_theta_crossing( return INFTY; const double s = -0.5 * c / b; - // Check if solution is in positive direction of flight and has correct sign + // Check if solution is in positive direction of flight and has correct + // sign if ((s > l) && (std::signbit(r.z + s * u.z) == sgn)) return s; @@ -1943,22 +2208,25 @@ extern "C" int openmc_mesh_bounding_box(int32_t index, double* ll, double* ur) return 0; } -extern "C" int openmc_mesh_material_volumes(int32_t index, int n_sample, - int bin, int result_size, void* result, int* hits, uint64_t* seed) +extern "C" int openmc_mesh_material_volumes(int32_t index, int nx, int ny, + int nz, int table_size, int32_t* materials, double* volumes) { - auto result_ = reinterpret_cast(result); - if (!result_) { - set_errmsg("Invalid result pointer passed to openmc_mesh_material_volumes"); - return OPENMC_E_INVALID_ARGUMENT; - } - if (int err = check_mesh(index)) return err; - int n = model::meshes[index]->material_volumes( - n_sample, bin, {result_, result_ + result_size}, seed); - *hits = n; - return (n == -1) ? OPENMC_E_ALLOCATE : 0; + try { + model::meshes[index]->material_volumes( + nx, ny, nz, table_size, materials, volumes); + } catch (const std::exception& e) { + set_errmsg(e.what()); + if (starts_with(e.what(), "Mesh")) { + return OPENMC_E_GEOMETRY; + } else { + return OPENMC_E_ALLOCATE; + } + } + + return 0; } extern "C" int openmc_mesh_get_plot_bins(int32_t index, Position origin, @@ -2394,7 +2662,8 @@ void MOABMesh::intersect_track(const moab::CartVect& start, moab::ErrorCode rval; vector tris; // get all intersections with triangles in the tet mesh - // (distances are relative to the start point, not the previous intersection) + // (distances are relative to the start point, not the previous + // intersection) rval = kdtree_->ray_intersect_triangles(kdtree_root_, FP_COINCIDENT, dir.array(), start.array(), tris, hits, 0, track_len); if (rval != moab::MB_SUCCESS) { @@ -2975,8 +3244,8 @@ void LibMesh::build_eqn_sys() void LibMesh::initialize() { if (!settings::libmesh_comm) { - fatal_error( - "Attempting to use an unstructured mesh without a libMesh communicator."); + fatal_error("Attempting to use an unstructured mesh without a libMesh " + "communicator."); } // assuming that unstructured meshes used in OpenMC are 3D @@ -3276,8 +3545,8 @@ void read_meshes(pugi::xml_node root) // Check to make sure multiple meshes in the same file don't share IDs int id = std::stoi(get_node_value(node, "id")); if (contains(mesh_ids, id)) { - fatal_error(fmt::format( - "Two or more meshes use the same unique ID '{}' in the same input file", + fatal_error(fmt::format("Two or more meshes use the same unique ID " + "'{}' in the same input file", id)); } mesh_ids.insert(id); diff --git a/src/particle.cpp b/src/particle.cpp index c63c84674..c51011d6b 100644 --- a/src/particle.cpp +++ b/src/particle.cpp @@ -293,7 +293,6 @@ void Particle::event_cross_surface() event() = TallyEvent::LATTICE; } else { // Particle crosses surface - // TODO: off-by-one const auto& surf {model::surfaces[surface_index()].get()}; // If BC, add particle to surface source before crossing surface if (surf->surf_source_ && surf->bc_) { diff --git a/tests/unit_tests/test_lib.py b/tests/unit_tests/test_lib.py index 64c16c238..8ab35335f 100644 --- a/tests/unit_tests/test_lib.py +++ b/tests/unit_tests/test_lib.py @@ -601,18 +601,18 @@ def test_regular_mesh(lib_init): mesh.set_parameters(lower_left=(-0.63, -0.63, -0.5), upper_right=(0.63, 0.63, 0.5)) vols = mesh.material_volumes() - assert len(vols) == 4 - for elem_vols in vols: + assert vols.num_elements == 4 + for i in range(vols.num_elements): + elem_vols = vols.by_element(i) assert sum(f[1] for f in elem_vols) == pytest.approx(1.26 * 1.26 / 4) - # If the mesh extends beyond the boundaries of the model, the volumes should - # still be reported correctly + # If the mesh extends beyond the boundaries of the model, we should get a + # GeometryError mesh.dimension = (1, 1, 1) mesh.set_parameters(lower_left=(-1.0, -1.0, -0.5), upper_right=(1.0, 1.0, 0.5)) - vols = mesh.material_volumes(100_000) - for elem_vols in vols: - assert sum(f[1] for f in elem_vols) == pytest.approx(1.26 * 1.26, 1e-2) + with pytest.raises(exc.GeometryError, match="not fully contained"): + vols = mesh.material_volumes() def test_regular_mesh_get_plot_bins(lib_init): @@ -683,11 +683,11 @@ def test_rectilinear_mesh(lib_init): mesh.set_grid([-w/2, -w/4, w/2], [-w/2, -w/4, w/2], [-0.5, 0.5]) vols = mesh.material_volumes() - assert len(vols) == 4 - assert sum(f[1] for f in vols[0]) == pytest.approx(w/4 * w/4) - assert sum(f[1] for f in vols[1]) == pytest.approx(w/4 * 3*w/4) - assert sum(f[1] for f in vols[2]) == pytest.approx(3*w/4 * w/4) - assert sum(f[1] for f in vols[3]) == pytest.approx(3*w/4 * 3*w/4) + assert vols.num_elements == 4 + assert sum(f[1] for f in vols.by_element(0)) == pytest.approx(w/4 * w/4) + assert sum(f[1] for f in vols.by_element(1)) == pytest.approx(w/4 * 3*w/4) + assert sum(f[1] for f in vols.by_element(2)) == pytest.approx(3*w/4 * w/4) + assert sum(f[1] for f in vols.by_element(3)) == pytest.approx(3*w/4 * 3*w/4) def test_cylindrical_mesh(lib_init): @@ -737,11 +737,11 @@ def test_cylindrical_mesh(lib_init): mesh.set_grid(r_grid, phi_grid, z_grid) vols = mesh.material_volumes() - assert len(vols) == 6 + assert vols.num_elements == 6 for i in range(0, 6, 2): - assert sum(f[1] for f in vols[i]) == pytest.approx(pi * 0.25**2 / 3) + assert sum(f[1] for f in vols.by_element(i)) == pytest.approx(pi * 0.25**2 / 3) for i in range(1, 6, 2): - assert sum(f[1] for f in vols[i]) == pytest.approx(pi * (0.5**2 - 0.25**2) / 3) + assert sum(f[1] for f in vols.by_element(i)) == pytest.approx(pi * (0.5**2 - 0.25**2) / 3) def test_spherical_mesh(lib_init): @@ -795,14 +795,14 @@ def test_spherical_mesh(lib_init): mesh.set_grid(r_grid, theta_grid, phi_grid) vols = mesh.material_volumes() - assert len(vols) == 12 + assert vols.num_elements == 12 d_theta = theta_grid[1] - theta_grid[0] d_phi = phi_grid[1] - phi_grid[0] for i in range(0, 12, 2): - assert sum(f[1] for f in vols[i]) == pytest.approx( + assert sum(f[1] for f in vols.by_element(i)) == pytest.approx( 0.25**3 / 3 * d_theta * d_phi * 2/pi) for i in range(1, 12, 2): - assert sum(f[1] for f in vols[i]) == pytest.approx( + assert sum(f[1] for f in vols.by_element(i)) == pytest.approx( (0.5**3 - 0.25**3) / 3 * d_theta * d_phi * 2/pi) diff --git a/tests/unit_tests/test_mesh.py b/tests/unit_tests/test_mesh.py index 3f307dda8..f1f5e9f67 100644 --- a/tests/unit_tests/test_mesh.py +++ b/tests/unit_tests/test_mesh.py @@ -1,4 +1,5 @@ from math import pi +from tempfile import TemporaryDirectory from pathlib import Path import numpy as np @@ -505,7 +506,7 @@ def test_mesh_get_homogenized_materials(): mesh.lower_left = (-1., -1., -1.) mesh.upper_right = (1., 1., 1.) mesh.dimension = (3, 1, 1) - m1, m2, m3 = mesh.get_homogenized_materials(model, n_samples=1_000_000) + m1, m2, m3 = mesh.get_homogenized_materials(model, n_samples=10_000) # Left mesh element should be only Fe56 assert m1.get_mass_density('Fe56') == pytest.approx(5.0) @@ -521,7 +522,7 @@ def test_mesh_get_homogenized_materials(): mesh_void.lower_left = (0.5, 0.5, -1.) mesh_void.upper_right = (1.5, 1.5, 1.) mesh_void.dimension = (1, 1, 1) - m4, = mesh_void.get_homogenized_materials(model, n_samples=1_000_000) + m4, = mesh_void.get_homogenized_materials(model, n_samples=(100, 100, 0)) # Mesh element that overlaps void should have half density assert m4.get_mass_density('H1') == pytest.approx(0.5, rel=1e-2) @@ -531,3 +532,86 @@ def test_mesh_get_homogenized_materials(): m5, = mesh_void.get_homogenized_materials( model, n_samples=1000, include_void=False) assert m5.get_mass_density('H1') == pytest.approx(1.0) + + +@pytest.fixture +def sphere_model(): + # Model with three materials separated by planes x=0 and z=0 + mats = [] + for i in range(3): + mat = openmc.Material() + mat.add_nuclide('H1', 1.0) + mat.set_density('g/cm3', float(i + 1)) + mats.append(mat) + + sph = openmc.Sphere(r=25.0, boundary_type='vacuum') + x0 = openmc.XPlane(0.0) + z0 = openmc.ZPlane(0.0) + cell1 = openmc.Cell(fill=mats[0], region=-sph & +x0 & +z0) + cell2 = openmc.Cell(fill=mats[1], region=-sph & -x0 & +z0) + cell3 = openmc.Cell(fill=mats[2], region=-sph & -z0) + model = openmc.Model() + model.geometry = openmc.Geometry([cell1, cell2, cell3]) + model.materials = openmc.Materials(mats) + return model + + +@pytest.mark.parametrize("n_rays", [1000, (10, 10, 0), (10, 0, 10), (0, 10, 10)]) +def test_material_volumes_regular_mesh(sphere_model, n_rays): + """Test the material_volumes method on a regular mesh""" + mesh = openmc.RegularMesh() + mesh.lower_left = (-1., -1., -1.) + mesh.upper_right = (1., 1., 1.) + mesh.dimension = (2, 2, 2) + volumes = mesh.material_volumes(sphere_model, n_rays) + mats = sphere_model.materials + np.testing.assert_almost_equal(volumes[mats[0].id], [0., 0., 0., 0., 0., 1., 0., 1.]) + np.testing.assert_almost_equal(volumes[mats[1].id], [0., 0., 0., 0., 1., 0., 1., 0.]) + np.testing.assert_almost_equal(volumes[mats[2].id], [1., 1., 1., 1., 0., 0., 0., 0.]) + assert volumes.by_element(4) == [(mats[1].id, 1.)] + assert volumes.by_element(0) == [(mats[2].id, 1.)] + + +def test_material_volumes_cylindrical_mesh(sphere_model): + """Test the material_volumes method on a cylindrical mesh""" + cyl_mesh = openmc.CylindricalMesh( + [0., 1.], [-1., 0., 1.,], [0.0, pi/4, 3*pi/4, 5*pi/4, 7*pi/4, 2*pi]) + volumes = cyl_mesh.material_volumes(sphere_model, (0, 100, 100)) + mats = sphere_model.materials + np.testing.assert_almost_equal(volumes[mats[0].id], [ + 0., 0., 0., 0., 0., + pi/8, pi/8, 0., pi/8, pi/8 + ]) + np.testing.assert_almost_equal(volumes[mats[1].id], [ + 0., 0., 0., 0., 0., + 0., pi/8, pi/4, pi/8, 0. + ]) + np.testing.assert_almost_equal(volumes[mats[2].id], [ + pi/8, pi/4, pi/4, pi/4, pi/8, + 0., 0., 0., 0., 0. + ]) + + +def test_mesh_material_volumes_serialize(): + materials = np.array([ + [1, -1, -2], + [-1, -2, -2], + [2, 1, -2], + [2, -2, -2] + ]) + volumes = np.array([ + [0.5, 0.5, 0.0], + [1.0, 0.0, 0.0], + [0.5, 0.5, 0.0], + [1.0, 0.0, 0.0] + ]) + volumes = openmc.MeshMaterialVolumes(materials, volumes) + with TemporaryDirectory() as tmpdir: + path = f'{tmpdir}/volumes.npz' + volumes.save(path) + new_volumes = openmc.MeshMaterialVolumes.from_npz(path) + + assert new_volumes.by_element(0) == [(1, 0.5), (None, 0.5)] + assert new_volumes.by_element(1) == [(None, 1.0)] + assert new_volumes.by_element(2) == [(2, 0.5), (1, 0.5)] + assert new_volumes.by_element(3) == [(2, 1.0)]