diff --git a/openmc/filter.py b/openmc/filter.py index 31d72b0dd4c..c58caeafcc8 100644 --- a/openmc/filter.py +++ b/openmc/filter.py @@ -86,7 +86,7 @@ def __new__(cls, name, bases, namespace, **kwargs): def _repeat_and_tile(bins, repeat_factor, data_size): - filter_bins = np.repeat(bins, repeat_factor) + filter_bins = np.repeat(bins, repeat_factor, axis=0) tile_factor = data_size // len(filter_bins) return np.tile(filter_bins, tile_factor) @@ -979,22 +979,13 @@ def get_pandas_dataframe(self, data_size, stride, **kwargs): Tally.get_pandas_dataframe(), CrossFilter.get_pandas_dataframe() """ - # Initialize dictionary to build Pandas Multi-index column - filter_dict = {} - # Append mesh ID as outermost index of multi-index mesh_key = f'mesh {self.mesh.id}' - # Determine index base (0-based for unstructured, 1-based otherwise) - 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): - filter_dict[mesh_key, label] = _repeat_and_tile( - np.arange(idx_start, idx_start + dim_size), stride, data_size) - stride *= dim_size - - return pd.DataFrame(filter_dict) + # Take the element indices from the mesh itself. + columns = [(mesh_key, label) for label in self.mesh.axis_labels] + indices = _repeat_and_tile(list(self.mesh.indices), stride, data_size) + return pd.DataFrame(indices, columns=columns) def to_xml_element(self): """Return XML Element representing the Filter. @@ -1318,36 +1309,23 @@ def get_pandas_dataframe(self, data_size, stride, **kwargs): Tally.get_pandas_dataframe(), CrossFilter.get_pandas_dataframe() """ - # Initialize Pandas DataFrame - df = pd.DataFrame() - - # Initialize dictionary to build Pandas Multi-index column - filter_dict = {} - # Append mesh ID as outermost index of multi-index mesh_key = f'mesh {self.mesh.id}' - # Number of surface-crossing bins per mesh element - dims = self.mesh.dimension - n_surfs = 4 * len(dims) - # Surface-crossing names derived from the mesh's axis labels current_names = _mesh_current_names(self.mesh) - - # 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 + n_surfs = len(current_names) + + # Take the element indices from the mesh itself, repeated once per + # surface-crossing bin, then append the surface column + columns = [(mesh_key, label) for label in self.mesh.axis_labels] + indices = _repeat_and_tile( + list(self.mesh.indices), stride * n_surfs, data_size) + filter_dict = dict(zip(columns, indices.T)) 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)]) + return pd.DataFrame(filter_dict) class CollisionFilter(Filter): diff --git a/tests/unit_tests/test_mesh_filter_dataframe.py b/tests/unit_tests/test_mesh_filter_dataframe.py new file mode 100644 index 00000000000..add8e6eabf3 --- /dev/null +++ b/tests/unit_tests/test_mesh_filter_dataframe.py @@ -0,0 +1,156 @@ +"""Mesh filter dataframes must match the mesh's own element indices. + +`MeshFilter` and `MeshSurfaceFilter` build their index columns from +`mesh.indices` rather than reconstructing them per axis. These tests pin the +column names, the ordering, and the index base for every mesh type, so that a +change to `indices` cannot silently alter every user's dataframe. +""" + +import numpy as np +import openmc +import pytest + + +@pytest.fixture +def regular_3d(): + mesh = openmc.RegularMesh() + mesh.lower_left = (-1.0, -1.0, -1.0) + mesh.upper_right = (1.0, 1.0, 1.0) + mesh.dimension = (2, 3, 4) + return mesh + + +@pytest.fixture +def regular_2d(): + mesh = openmc.RegularMesh() + mesh.lower_left = (-1.0, -1.0) + mesh.upper_right = (1.0, 1.0) + mesh.dimension = (3, 2) + return mesh + + +@pytest.fixture +def rectilinear(): + mesh = openmc.RectilinearMesh() + mesh.x_grid = [-1.0, 0.0, 1.0] + mesh.y_grid = [-1.0, 0.0, 2.0] + mesh.z_grid = [-1.0, 1.0] + return mesh + + +@pytest.fixture +def cylindrical(): + return openmc.CylindricalMesh( + r_grid=[0.0, 1.0, 2.0], + phi_grid=[0.0, np.pi, 2 * np.pi], + z_grid=[-1.0, 0.0, 1.0], + ) + + +@pytest.fixture +def spherical(): + return openmc.SphericalMesh( + r_grid=[0.0, 1.0, 2.0], + theta_grid=[0.0, np.pi / 2, np.pi], + phi_grid=[0.0, np.pi, 2 * np.pi], + ) + + +ALL_MESHES = ('regular_3d', 'regular_2d', 'rectilinear', 'cylindrical', + 'spherical') + +EXPECTED_LABELS = { + 'regular_3d': ('x', 'y', 'z'), + 'regular_2d': ('x', 'y'), + 'rectilinear': ('x', 'y', 'z'), + 'cylindrical': ('r', 'phi', 'z'), + 'spherical': ('r', 'theta', 'phi'), +} + + +@pytest.mark.parametrize('mesh_name', ALL_MESHES) +def test_mesh_filter_columns(mesh_name, request): + """Column names come from the mesh's own axis labels.""" + mesh = request.getfixturevalue(mesh_name) + filt = openmc.MeshFilter(mesh) + df = filt.get_pandas_dataframe(filt.num_bins, 1) + + key = f'mesh {mesh.id}' + assert list(df.columns) == [(key, ax) for ax in EXPECTED_LABELS[mesh_name]] + assert len(df) == filt.num_bins + + +@pytest.mark.parametrize('mesh_name', ALL_MESHES) +def test_mesh_filter_matches_mesh_indices(mesh_name, request): + """Rows are exactly the mesh's element indices, in bin order. + + This is the property the implementation relies on. Structured meshes report + 1-based indices, so the first row is all ones and the last row is the + dimension tuple. + """ + mesh = request.getfixturevalue(mesh_name) + filt = openmc.MeshFilter(mesh) + df = filt.get_pandas_dataframe(filt.num_bins, 1) + + expected = list(mesh.indices) + assert [tuple(row) for row in df.to_numpy()] == [tuple(i) for i in expected] + + # Structured meshes are 1-based + assert tuple(df.to_numpy()[0]) == tuple([1] * len(EXPECTED_LABELS[mesh_name])) + assert tuple(df.to_numpy()[-1]) == tuple(mesh.dimension) + + +@pytest.mark.parametrize('mesh_name', ALL_MESHES) +def test_mesh_surface_filter_columns(mesh_name, request): + """Surface filter adds a 'surf' column after the axis columns.""" + mesh = request.getfixturevalue(mesh_name) + filt = openmc.MeshSurfaceFilter(mesh) + df = filt.get_pandas_dataframe(filt.num_bins, 1) + + key = f'mesh {mesh.id}' + labels = EXPECTED_LABELS[mesh_name] + assert list(df.columns) == [(key, ax) for ax in labels] + [(key, 'surf')] + assert len(df) == filt.num_bins + + +@pytest.mark.parametrize('mesh_name', ALL_MESHES) +def test_mesh_surface_filter_bin_order(mesh_name, request): + """Surface names cycle fastest, element indices slowest.""" + mesh = request.getfixturevalue(mesh_name) + filt = openmc.MeshSurfaceFilter(mesh) + df = filt.get_pandas_dataframe(filt.num_bins, 1) + + key = f'mesh {mesh.id}' + labels = EXPECTED_LABELS[mesh_name] + n_surfs = 4 * len(labels) + + surfs = list(df[(key, 'surf')]) + assert surfs[:4] == [f'{labels[0]}-min out', f'{labels[0]}-min in', + f'{labels[0]}-max out', f'{labels[0]}-max in'] + # The surface column repeats with period n_surfs + assert surfs[:n_surfs] == surfs[n_surfs:2 * n_surfs] + + # The element index is constant across one element's surface bins + first = df.iloc[:n_surfs][[(key, ax) for ax in labels]].to_numpy() + assert (first == first[0]).all() + + +def test_unstructured_mesh_filter_is_zero_based(): + """Unstructured meshes stay 0-based with a single element_index column. + + This is the case the removed `idx_start` special case in MeshFilter existed + for. UnstructuredMesh.indices is already 0-based, so delegating to it gives + the same result without the isinstance check. + """ + mesh = openmc.UnstructuredMesh('dummy.exo', 'moab') + # Stand in for data that would normally come from a statepoint + mesh._has_statepoint_data = True + mesh.n_elements = 4 + mesh._volumes = np.ones(4) + + filt = openmc.MeshFilter(mesh) + df = filt.get_pandas_dataframe(filt.num_bins, 1) + + key = f'mesh {mesh.id}' + assert list(df.columns) == [(key, 'element_index')] + assert [tuple(row) for row in df.to_numpy()] == [(0,), (1,), (2,), (3,)]