diff --git a/openmc/filter.py b/openmc/filter.py index 87aeb70c3a..b16f485df9 100644 --- a/openmc/filter.py +++ b/openmc/filter.py @@ -30,12 +30,19 @@ _FILTER_TYPES = ( 'weight', 'meshborn', 'meshsurface', 'meshmaterial', 'reaction', ) -_CURRENT_NAMES = ( - 'x-min out', 'x-min in', 'x-max out', 'x-max in', - 'y-min out', 'y-min in', 'y-max out', 'y-max in', - 'z-min out', 'z-min in', 'z-max out', 'z-max in' -) +def _mesh_current_names(mesh): + """Return surface-crossing bin names derived from a mesh's axis labels. + For each axis (in dimension order) the four crossing names are produced as + "-min out", "-min in", "-max out", "-max in". This + yields x/y/z names for Cartesian meshes, r/phi/z for cylindrical, and + r/theta/phi for spherical. + """ + names = [] + for axis in mesh.axis_labels: + names += [f'{axis}-min out', f'{axis}-min in', + f'{axis}-max out', f'{axis}-max in'] + return names class FilterMeta(ABCMeta): @@ -996,7 +1003,7 @@ class MeshFilter(Filter): idx_start = 0 if isinstance(self.mesh, openmc.UnstructuredMesh) else 1 # Generate a multi-index sub-column for each axis - for label, dim_size in zip(self.mesh._axis_labels, self.mesh.dimension): + for label, dim_size in zip(self.mesh.axis_labels, self.mesh.dimension): filter_dict[mesh_key, label] = _repeat_and_tile( np.arange(idx_start, idx_start + dim_size), stride, data_size) stride *= dim_size @@ -1273,7 +1280,8 @@ class MeshSurfaceFilter(MeshFilter): Unique identifier for the filter bins : list of tuple A list of mesh indices / surfaces for each filter bin, e.g. [(1, 1, - 'x-min out'), (1, 1, 'x-min in'), ...] + 'x-min out'), (1, 1, 'x-min in'), ...]. Surface names use the mesh's + axis labels (e.g. r/phi/z for a cylindrical mesh). num_bins : Integral The number of filter bins @@ -1287,10 +1295,12 @@ class MeshSurfaceFilter(MeshFilter): cv.check_type('filter mesh', mesh, openmc.MeshBase) self._mesh = mesh - # Take the product of mesh indices and current names - n_dim = mesh.n_dimension + # Take the product of mesh indices and surface-crossing names, using + # the mesh's own axis labels so cylindrical/spherical meshes get the + # correct names. + current_names = _mesh_current_names(mesh) self.bins = [mesh_tuple + (surf,) for mesh_tuple, surf in - product(mesh.indices, _CURRENT_NAMES[:4*n_dim])] + product(mesh.indices, current_names)] def get_pandas_dataframe(self, data_size, stride, **kwargs): """Builds a Pandas DataFrame for the Filter's bins. @@ -1309,9 +1319,11 @@ class MeshSurfaceFilter(MeshFilter): Returns ------- pandas.DataFrame - A Pandas DataFrame with three columns describing the x,y,z mesh - cell indices corresponding to each filter bin. The number of rows - in the DataFrame is the same as the total number of bins in the + A Pandas DataFrame with columns describing the mesh cell indices + corresponding to each filter bin, plus a surface column. Column + names depend on the mesh type (e.g., x/y/z for RegularMesh, r/phi/z + for CylindricalMesh, r/theta/phi for SphericalMesh). The number of + rows in the DataFrame is the same as the total number of bins in the corresponding tally, with the filter bin appropriately tiled to map to the corresponding tally bins. @@ -1329,34 +1341,24 @@ class MeshSurfaceFilter(MeshFilter): # Append mesh ID as outermost index of multi-index mesh_key = f'mesh {self.mesh.id}' - # Find mesh dimensions - use 3D indices for simplicity - n_surfs = 4 * len(self.mesh.dimension) - if len(self.mesh.dimension) == 3: - nx, ny, nz = self.mesh.dimension - elif len(self.mesh.dimension) == 2: - nx, ny = self.mesh.dimension - nz = 1 - else: - nx = self.mesh.dimension - ny = nz = 1 + # Number of surface-crossing bins per mesh element + dims = self.mesh.dimension + n_surfs = 4 * len(dims) - # Generate multi-index sub-column for x-axis - filter_dict[mesh_key, 'x'] = _repeat_and_tile( - np.arange(1, nx + 1), n_surfs * stride, data_size) + # Surface-crossing names derived from the mesh's axis labels + current_names = _mesh_current_names(self.mesh) - # Generate multi-index sub-column for y-axis - if len(self.mesh.dimension) > 1: - filter_dict[mesh_key, 'y'] = _repeat_and_tile( - np.arange(1, ny + 1), n_surfs * nx * stride, data_size) - - # Generate multi-index sub-column for z-axis - if len(self.mesh.dimension) > 2: - filter_dict[mesh_key, 'z'] = _repeat_and_tile( - np.arange(1, nz + 1), n_surfs * nx * ny * stride, data_size) + # Generate a multi-index sub-column for each axis, using the mesh's own + # axis labels so curvilinear meshes are labeled correctly. + axis_stride = n_surfs * stride + for label, dim_size in zip(self.mesh.axis_labels, dims): + filter_dict[mesh_key, label] = _repeat_and_tile( + np.arange(1, dim_size + 1), axis_stride, data_size) + axis_stride *= dim_size # Generate multi-index sub-column for surface filter_dict[mesh_key, 'surf'] = _repeat_and_tile( - _CURRENT_NAMES[:n_surfs], stride, data_size) + current_names, stride, data_size) # Initialize a Pandas DataFrame from the mesh dictionary return pd.concat([df, pd.DataFrame(filter_dict)]) diff --git a/openmc/mesh.py b/openmc/mesh.py index ff8e145772..bcb3ad2538 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -259,6 +259,12 @@ class MeshBase(IDManagerMixin, ABC): def n_elements(self): pass + @property + @abstractmethod + def axis_labels(self): + """tuple of str : Names of the mesh axes, one per dimension.""" + pass + def __repr__(self): string = type(self).__name__ + '\n' string += '{0: <16}{1}{2}\n'.format('\tID', '=\t', self._id) @@ -538,7 +544,8 @@ class StructuredMesh(MeshBase): @property @abstractmethod - def _axis_labels(self): + def axis_labels(self): + """tuple of str : Names of the mesh axes, one per dimension.""" pass @property @@ -1042,6 +1049,9 @@ class RegularMesh(StructuredMesh): The number of mesh cells in each direction (x, y, z). n_dimension : int Number of mesh dimensions. + axis_labels : tuple of str + Names of the mesh axes ('x', 'y', 'z'), truncated to the mesh + dimensionality. lower_left : Iterable of float 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. @@ -1085,7 +1095,7 @@ class RegularMesh(StructuredMesh): return None @property - def _axis_labels(self): + def axis_labels(self): return ('x', 'y', 'z')[:self.n_dimension] @property @@ -1569,6 +1579,8 @@ class RectilinearMesh(StructuredMesh): The number of mesh cells in each direction (x, y, z). n_dimension : int Number of mesh dimensions (always 3 for a RectilinearMesh). + axis_labels : tuple of str + Names of the mesh axes ('x', 'y', 'z'). x_grid : numpy.ndarray 1-D array of mesh boundary points along the x-axis. y_grid : numpy.ndarray @@ -1602,7 +1614,7 @@ class RectilinearMesh(StructuredMesh): return 3 @property - def _axis_labels(self): + def axis_labels(self): return ('x', 'y', 'z') @property @@ -1870,6 +1882,8 @@ class CylindricalMesh(StructuredMesh): The number of mesh cells in each direction (r_grid, phi_grid, z_grid). n_dimension : int Number of mesh dimensions (always 3 for a CylindricalMesh). + axis_labels : tuple of str + Names of the mesh axes ('r', 'phi', 'z'). r_grid : numpy.ndarray 1-D array of mesh boundary points along the r-axis. Requirement is r >= 0. @@ -1924,7 +1938,7 @@ class CylindricalMesh(StructuredMesh): return 3 @property - def _axis_labels(self): + def axis_labels(self): return ('r', 'phi', 'z') @property @@ -2307,6 +2321,8 @@ class SphericalMesh(StructuredMesh): theta_grid, phi_grid). n_dimension : int Number of mesh dimensions (always 3 for a SphericalMesh). + axis_labels : tuple of str + Names of the mesh axes ('r', 'theta', 'phi'). r_grid : numpy.ndarray 1-D array of mesh boundary points along the r-axis. Requirement is r >= 0. @@ -2361,7 +2377,7 @@ class SphericalMesh(StructuredMesh): return 3 @property - def _axis_labels(self): + def axis_labels(self): return ('r', 'theta', 'phi') @property @@ -2949,7 +2965,7 @@ class UnstructuredMesh(MeshBase): return 3 @property - def _axis_labels(self): + def axis_labels(self): return ('element_index',) @property