MeshSurfaceFilter partial current labels (#4021)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

include/openmc/mesh.h

@@ -236,6 +236,12 @@ public:
   //! \param[in] bin Mesh bin to generate a label for
   virtual std::string bin_label(int bin) const = 0;
 
+  //! Axis names (per dimension) used when labeling surface tally bins
+  virtual std::array<const char*, 3> axis_labels() const;
+
+  //! Build the surface component of a mesh surface tally bin label
+  std::string surface_bin_label(int surf_index) const;
+
   //! Get the volume of a mesh bin
   //
   //! \param[in] bin Bin to return the volume for
@@ -588,6 +594,8 @@ public:
 
   static const std::string mesh_type;
 
+  std::array<const char*, 3> axis_labels() const override;
+
   Position sample_element(const MeshIndex& ijk, uint64_t* seed) const override;
 
   MeshDistance distance_to_grid_boundary(const MeshIndex& ijk, int i,
@@ -653,6 +661,8 @@ public:
 
   static const std::string mesh_type;
 
+  std::array<const char*, 3> axis_labels() const override;
+
   Position sample_element(const MeshIndex& ijk, uint64_t* seed) const override;
 
   MeshDistance distance_to_grid_boundary(const MeshIndex& ijk, int i,

include/openmc/tallies/filter_meshsurface.h

@@ -22,21 +22,6 @@ public:
   // Accessors
 
   void set_mesh(int32_t mesh) override;
-
-  enum class MeshDir {
-    OUT_LEFT,   // x min
-    IN_LEFT,    // x min
-    OUT_RIGHT,  // x max
-    IN_RIGHT,   // x max
-    OUT_BACK,   // y min
-    IN_BACK,    // y min
-    OUT_FRONT,  // y max
-    IN_FRONT,   // y max
-    OUT_BOTTOM, // z min
-    IN_BOTTOM,  // z min
-    OUT_TOP,    // z max
-    IN_TOP      // z max
-  };
 };
 
 } // namespace openmc

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
+    "<axis>-min out", "<axis>-min in", "<axis>-max out", "<axis>-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[:n_surfs], stride, data_size)
 
         # Initialize a Pandas DataFrame from the mesh dictionary
         return pd.concat([df, pd.DataFrame(filter_dict)])

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

openmc/mgxs/mdgxs.py

@@ -877,7 +877,7 @@ class MDGXS(MGXS):
         # energy groups such that data is from fast to thermal
         if self.domain_type == 'mesh':
             mesh_str = f'mesh {self.domain.id}'
-            mesh_cols = [(mesh_str, label) for label in self.domain._axis_labels]
+            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,7 +2134,7 @@ class MGXS:
         # energy groups such that data is from fast to thermal
         if self.domain_type == 'mesh':
             mesh_str = f'mesh {self.domain.id}'
-            mesh_cols = [(mesh_str, label) for label in self.domain._axis_labels]
+            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)

src/mesh.cpp

@@ -429,6 +429,25 @@ vector<double> Mesh::volumes() const
   return volumes;
 }
 
+//! Default (Cartesian) axis labels used for surface bin labels.
+std::array<const char*, 3> Mesh::axis_labels() const
+{
+  return {"x", "y", "z"};
+}
+
+//! Build the surface component of a mesh surface tally bin label.
+//! surf_index: 0=out/min, 1=in/min, 2=out/max, 3=in/max
+std::string Mesh::surface_bin_label(int surf_index) const
+{
+  auto labels = this->axis_labels();
+  int dim = surf_index / 4;
+  int code = surf_index % 4;
+  bool incoming = (code == 1) || (code == 3);
+  bool max = (code == 2) || (code == 3);
+  return fmt::format(" {}, {}-{}", incoming ? "Incoming" : "Outgoing",
+    labels[dim], max ? "max" : "min");
+}
+
 void Mesh::material_volumes(int nx, int ny, int nz, int table_size,
   int32_t* materials, double* volumes) const
 {
@@ -1814,6 +1833,11 @@ std::string CylindricalMesh::get_mesh_type() const
   return mesh_type;
 }
 
+std::array<const char*, 3> CylindricalMesh::axis_labels() const
+{
+  return {"r", "phi", "z"};
+}
+
 StructuredMesh::MeshIndex CylindricalMesh::get_indices(
   Position r, bool& in_mesh) const
 {
@@ -2107,6 +2131,11 @@ std::string SphericalMesh::get_mesh_type() const
   return mesh_type;
 }
 
+std::array<const char*, 3> SphericalMesh::axis_labels() const
+{
+  return {"r", "theta", "phi"};
+}
+
 StructuredMesh::MeshIndex SphericalMesh::get_indices(
   Position r, bool& in_mesh) const
 {

src/tallies/filter_meshsurface.cpp

@@ -28,52 +28,14 @@ std::string MeshSurfaceFilter::text_label(int bin) const
   auto& mesh = *model::meshes[mesh_];
   int n_dim = mesh.n_dimension_;
 
-  // Get flattend mesh index and surface index.
+  // Get flattened mesh index and surface index.
   int i_mesh = bin / (4 * n_dim);
-  MeshDir surf_dir = static_cast<MeshDir>(bin % (4 * n_dim));
+  int surf_index = bin % (4 * n_dim);
 
-  // Get mesh index part of label.
+  // Get mesh index part of label, then append the surface part.
+  // The surface is labeled by the underlying mesh.
   std::string out = MeshFilter::text_label(i_mesh);
-
-  // Get surface part of label.
-  switch (surf_dir) {
-  case MeshDir::OUT_LEFT:
-    out += " Outgoing, x-min";
-    break;
-  case MeshDir::IN_LEFT:
-    out += " Incoming, x-min";
-    break;
-  case MeshDir::OUT_RIGHT:
-    out += " Outgoing, x-max";
-    break;
-  case MeshDir::IN_RIGHT:
-    out += " Incoming, x-max";
-    break;
-  case MeshDir::OUT_BACK:
-    out += " Outgoing, y-min";
-    break;
-  case MeshDir::IN_BACK:
-    out += " Incoming, y-min";
-    break;
-  case MeshDir::OUT_FRONT:
-    out += " Outgoing, y-max";
-    break;
-  case MeshDir::IN_FRONT:
-    out += " Incoming, y-max";
-    break;
-  case MeshDir::OUT_BOTTOM:
-    out += " Outgoing, z-min";
-    break;
-  case MeshDir::IN_BOTTOM:
-    out += " Incoming, z-min";
-    break;
-  case MeshDir::OUT_TOP:
-    out += " Outgoing, z-max";
-    break;
-  case MeshDir::IN_TOP:
-    out += " Incoming, z-max";
-    break;
-  }
+  out += mesh.surface_bin_label(surf_index);
 
   return out;
 }

tests/unit_tests/test_filter_meshsurface.py

@@ -0,0 +1,130 @@
+from math import pi
+from pathlib import Path
+
+import numpy as np
+import openmc
+import pytest
+
+
+@pytest.mark.parametrize('mesh_type', ('regular', 'rectilinear', 'cylindrical', 'spherical'))
+def test_axis_labels(mesh_type):
+    if mesh_type == 'regular':
+        mesh = openmc.RegularMesh()
+        mesh.lower_left = np.asarray([0.0]*3)
+        mesh.width = np.asarray([0.0]*3)
+        mesh.dimension = (1, 1, 1)
+    elif mesh_type == 'rectilinear':
+        mesh = openmc.RectilinearMesh()
+        mesh.x_grid = np.asarray([0.0, 1.0])
+        mesh.y_grid = np.asarray([0.0, 1.0])
+        mesh.z_grid = np.asarray([0.0, 1.0])
+    elif mesh_type == 'cylindrical':
+        r_grid = np.asarray([0.0, 1.0])
+        z_grid = np.asarray([0.0, 1.0])
+        p_grid = np.asarray([0.0, 1.0])
+        mesh = openmc.CylindricalMesh(r_grid=r_grid, z_grid=z_grid, phi_grid=p_grid)
+    elif mesh_type == 'spherical':
+        r_grid = np.asarray([0.0, 1.0])
+        t_grid = np.asarray([0.0, 1.0])
+        p_grid = np.asarray([0.0, 1.0])
+        mesh = openmc.SphericalMesh(r_grid=r_grid, theta_grid=t_grid, phi_grid=p_grid)
+    else:
+        raise ValueError(mesh_type)
+
+    filt = openmc.MeshSurfaceFilter(mesh)
+    assert len(filt.bins) == 12
+    bin_names = [b[3] for b in filt.bins]
+    if mesh_type in {'regular', 'rectilinear'}:
+        assert bin_names[:8] == ['x-min out', 'x-min in', 'x-max out', 'x-max in',
+                                 'y-min out', 'y-min in', 'y-max out', 'y-max in']
+    if mesh_type in {'regular', 'rectilinear', 'cylindrical'}:
+        assert bin_names[8:] == ['z-min out', 'z-min in', 'z-max out', 'z-max in']
+    if mesh_type in {'cylindrical', 'spherical'}:
+        assert bin_names[:4] == ['r-min out', 'r-min in', 'r-max out', 'r-max in']
+    if mesh_type == 'spherical':
+        assert bin_names[4:] == ['theta-min out', 'theta-min in', 'theta-max out', 'theta-max in',
+                                 'phi-min out', 'phi-min in', 'phi-max out', 'phi-max in']
+    if mesh_type == 'cylindrical':
+        assert bin_names[4:8] == ['phi-min out', 'phi-min in', 'phi-max out', 'phi-max in']
+
+
+
+def test_mesh_surface_labels(run_in_tmpdir):
+    """Check that mesh surface labels use the mesh coordinate system."""
+    sphere = openmc.Sphere(r=1.0, boundary_type='vacuum')
+    model = openmc.Model(geometry=openmc.Geometry([openmc.Cell(region=-sphere)]))
+    model.settings.run_mode = 'fixed source'
+    model.settings.batches = 1
+    model.settings.particles = 1
+    model.settings.source = openmc.IndependentSource(
+        space=openmc.stats.Point((0.1, 0.1, 0.1))
+    )
+
+    mesh_1d = openmc.RegularMesh()
+    mesh_1d.dimension = (2,)
+    mesh_1d.lower_left = (-1.0,)
+    mesh_1d.upper_right = (1.0,)
+
+    mesh_2d = openmc.RegularMesh()
+    mesh_2d.dimension = (2, 1)
+    mesh_2d.lower_left = (-1.0, -1.0)
+    mesh_2d.upper_right = (1.0, 1.0)
+
+    mesh_3d = openmc.RegularMesh()
+    mesh_3d.dimension = (2, 1, 1)
+    mesh_3d.lower_left = (-1.0, -1.0, -1.0)
+    mesh_3d.upper_right = (1.0, 1.0, 1.0)
+
+    rectilinear_mesh = openmc.RectilinearMesh()
+    rectilinear_mesh.x_grid = (-1.0, 0.0, 1.0)
+    rectilinear_mesh.y_grid = (-1.0, 1.0)
+    rectilinear_mesh.z_grid = (-1.0, 1.0)
+    cylindrical_mesh = openmc.CylindricalMesh(
+        r_grid=(0.0, 0.5, 1.0),
+        phi_grid=(0.0, 2*pi),
+        z_grid=(-1.0, 1.0),
+    )
+    spherical_mesh = openmc.SphericalMesh(
+        r_grid=(0.0, 0.5, 1.0),
+        theta_grid=(0.0, pi),
+        phi_grid=(0.0, 2*pi),
+    )
+
+    tally_meshes = {}
+    for mesh in (
+        mesh_1d, mesh_2d, mesh_3d, rectilinear_mesh, cylindrical_mesh,
+        spherical_mesh
+    ):
+        tally = openmc.Tally()
+        tally.filters = [openmc.MeshSurfaceFilter(mesh)]
+        tally.scores = ['current']
+        model.tallies.append(tally)
+        tally_meshes[tally.id] = mesh
+
+    model.run()
+    output = Path('tallies.out').read_text()
+
+    for tally_id, mesh in tally_meshes.items():
+        header = f'TALLY {tally_id}'
+        start = output.index(header) + len(header)
+        end = output.find('TALLY', start)
+        section = output[start:] if end == -1 else output[start:end]
+
+        n_dim = mesh.n_dimension
+        labels = [
+            line for line in section.splitlines() if 'Mesh Index' in line
+        ][:4*n_dim + 1]
+
+        expected = []
+        for axis in mesh.axis_labels:
+            expected.extend([
+                f'Outgoing, {axis}-min',
+                f'Incoming, {axis}-min',
+                f'Outgoing, {axis}-max',
+                f'Incoming, {axis}-max',
+            ])
+        expected.append(expected[0])
+
+        assert len(labels) == len(expected)
+        for line, label in zip(labels, expected):
+            assert label in line