Fix MeshFilter.get_pandas_dataframe to handle all mesh types (#3817)

Author: paulromano

openmc/filter.py

@@ -972,53 +972,36 @@ def get_pandas_dataframe(self, data_size, stride, **kwargs):
         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
-            corresponding tally, with the filter bin appropriately tiled to map
-            to the corresponding tally bins.
+            A Pandas DataFrame with columns describing the mesh cell indices
+            corresponding to each filter bin. Column names depend on the mesh
+            type (e.g., x/y/z for RegularMesh, r/phi/z for CylindricalMesh,
+            r/theta/phi for SphericalMesh, or element index for
+            UnstructuredMesh). 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.
 
         See also
         --------
         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}'
 
-        # Find mesh dimensions - use 3D indices for simplicity
-        n_dim = len(self.mesh.dimension)
-        if n_dim == 3:
-            nx, ny, nz = self.mesh.dimension
-        elif n_dim == 2:
-            nx, ny = self.mesh.dimension
-            nz = 1
-        else:
-            nx = self.mesh.dimension
-            ny = nz = 1
-
-        # Generate multi-index sub-column for x-axis
-        filter_dict[mesh_key, 'x'] = _repeat_and_tile(
-            np.arange(1, nx + 1), stride, data_size)
-
-        # Generate multi-index sub-column for y-axis
-        filter_dict[mesh_key, 'y'] = _repeat_and_tile(
-            np.arange(1, ny + 1), nx * stride, data_size)
-
-        # Generate multi-index sub-column for z-axis
-        filter_dict[mesh_key, 'z'] = _repeat_and_tile(
-            np.arange(1, nz + 1), nx * ny * stride, data_size)
+        # Determine index base (0-based for unstructured, 1-based otherwise)
+        idx_start = 0 if isinstance(self.mesh, openmc.UnstructuredMesh) else 1
 
-        # Initialize a Pandas DataFrame from the mesh dictionary
-        df = pd.concat([df, pd.DataFrame(filter_dict)])
+        # 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 df
+        return pd.DataFrame(filter_dict)
 
     def to_xml_element(self):
         """Return XML Element representing the Filter.

openmc/mesh.py

@@ -236,12 +236,12 @@ def name(self, name: str):
             self._name = name
         else:
             self._name = ''
-            
+
     @property
     @abstractmethod
     def lower_left(self):
         pass
-        
+
     @property
     @abstractmethod
     def upper_right(self):
@@ -255,7 +255,7 @@ def bounding_box(self) -> openmc.BoundingBox:
     @abstractmethod
     def indices(self):
         pass
-        
+
     @property
     @abstractmethod
     def n_elements(self):
@@ -537,6 +537,11 @@ def dimension(self):
     def n_dimension(self):
         pass
 
+    @property
+    @abstractmethod
+    def _axis_labels(self):
+        pass
+
     @property
     @abstractmethod
     def _grids(self):
@@ -636,7 +641,7 @@ def centroids(self):
         s0 = (slice(0, -1),)*ndim + (slice(None),)
         s1 = (slice(1, None),)*ndim + (slice(None),)
         return (vertices[s0] + vertices[s1]) / 2
-    
+
     @property
     def n_elements(self):
         return np.prod(self.dimension)
@@ -995,6 +1000,10 @@ def n_dimension(self):
         else:
             return None
 
+    @property
+    def _axis_labels(self):
+        return ('x', 'y', 'z')[:self.n_dimension]
+
     @property
     def lower_left(self):
         return self._lower_left
@@ -1475,6 +1484,10 @@ def dimension(self):
     def n_dimension(self):
         return 3
 
+    @property
+    def _axis_labels(self):
+        return ('x', 'y', 'z')
+
     @property
     def x_grid(self):
         return self._x_grid
@@ -1709,6 +1722,10 @@ def dimension(self):
     def n_dimension(self):
         return 3
 
+    @property
+    def _axis_labels(self):
+        return ('r', 'phi', 'z')
+
     @property
     def origin(self):
         return self._origin
@@ -2156,6 +2173,10 @@ def dimension(self):
     def n_dimension(self):
         return 3
 
+    @property
+    def _axis_labels(self):
+        return ('r', 'theta', 'phi')
+
     @property
     def origin(self):
         return self._origin
@@ -2671,6 +2692,10 @@ def dimension(self):
     def n_dimension(self):
         return 3
 
+    @property
+    def _axis_labels(self):
+        return ('element_index',)
+
     @property
     @require_statepoint_data
     def indices(self):

openmc/mgxs/mdgxs.py

@@ -877,8 +877,8 @@ def get_pandas_dataframe(self, groups='all', nuclides='all',
         # energy groups such that data is from fast to thermal
         if self.domain_type == 'mesh':
             mesh_str = f'mesh {self.domain.id}'
-            df.sort_values(by=[(mesh_str, 'x'), (mesh_str, 'y'),
-                               (mesh_str, 'z')] + columns, inplace=True)
+            mesh_cols = [(mesh_str, label) for label in self.domain._axis_labels]
+            df.sort_values(by=mesh_cols + columns, inplace=True)
         else:
             df.sort_values(by=[self.domain_type] + columns, inplace=True)
 

openmc/mgxs/mgxs.py

@@ -2134,8 +2134,8 @@ def get_pandas_dataframe(self, groups='all', nuclides='all',
         # energy groups such that data is from fast to thermal
         if self.domain_type == 'mesh':
             mesh_str = f'mesh {self.domain.id}'
-            df.sort_values(by=[(mesh_str, 'x'), (mesh_str, 'y'),
-                               (mesh_str, 'z')] + columns, inplace=True)
+            mesh_cols = [(mesh_str, label) for label in self.domain._axis_labels]
+            df.sort_values(by=mesh_cols + columns, inplace=True)
         else:
             df.sort_values(by=[self.domain_type] + columns, inplace=True)