Build mesh filter dataframe columns from mesh.indices

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

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,)]