Repository: openmc-dev/openmc

Issue:
Shuffle array shapes returned from mesh classes
Based on discussion with @paulromano and @eepeterson:

The shape of vertices, and centroids returned from our mesh classes is currently `(ndim, dim0, ..., dimn)`. While this is convenient for unpacking x, y, and z coordinates in a way that is useful for some plotting operations, 

```python
xx, yy, zz = mesh.vertices
```

it makes what indexing a centroid or vertex by an `ijk` tuple (i.e. `(2, 3, 6)`  counterintuitive

```python
centroid = mesh.centroids[(:, *ijk)]
```

These properties of the mesh class will be updated to a shape of `(dim0, ..., dimn, ndim)` to allow indexing to occur more naturally:

```python
centroid = mesh.centroids[ijk]
```

Note that x, y, z values can still be unpacked fairly easily with the following command:

```python
xx, yy, zz = np.rollaxis(mesh.vertices, -1)
```



Discussion:

Resolving patch:
diff --git a/openmc/mesh.py b/openmc/mesh.py
--- a/openmc/mesh.py
+++ b/openmc/mesh.py
@@ -179,18 +179,18 @@ def vertices(self):
         -------
         vertices : numpy.ndarray
             Returns a numpy.ndarray representing the coordinates of the mesh
-            vertices with a shape equal to (ndim, dim1 + 1, ..., dimn + 1).  Can be
-            unpacked along the first dimension with xx, yy, zz = mesh.vertices.
+            vertices with a shape equal to (dim1 + 1, ..., dimn + 1, ndim). X, Y, Z values
+            can be unpacked with xx, yy, zz = np.rollaxis(mesh.vertices, -1).
 
         """
         return self._generate_vertices(*self._grids)
 
     @staticmethod
     def _generate_vertices(i_grid, j_grid, k_grid):
-        """Returns an array with shape (3, i_grid.size+1, j_grid.size+1, k_grid.size+1)
+        """Returns an array with shape (i_grid.size, j_grid.size, k_grid.size, 3)
            containing the corner vertices of mesh elements.
         """
-        return np.stack(np.meshgrid(i_grid, j_grid, k_grid, indexing='ij'), axis=0)
+        return np.stack(np.meshgrid(i_grid, j_grid, k_grid, indexing='ij'), axis=-1)
 
     @staticmethod
     def _generate_edge_midpoints(grids):
@@ -206,7 +206,7 @@ def _generate_edge_midpoints(grids):
         midpoint_grids : list of numpy.ndarray
             The edge midpoints for the i, j, and k midpoints of each element in
             i, j, k ordering. The shapes of the resulting grids are
-            [(3, ni-1, nj, nk), (3, ni, nj-1, nk), (3, ni, nj, nk-1)]
+            [(ni-1, nj, nk, 3), (ni, nj-1, nk, 3), (ni, nj, nk-1, 3)]
         """
         # generate a set of edge midpoints for each dimension
         midpoint_grids = []
@@ -216,7 +216,7 @@ def _generate_edge_midpoints(grids):
         # each grid is comprised of the mid points for one dimension and the
         # corner vertices of the other two
         for dims in ((0, 1, 2), (1, 0, 2), (2, 0, 1)):
-            # compute the midpoints along the first dimension
+            # compute the midpoints along the last dimension
             midpoints = grids[dims[0]][:-1] + 0.5 * np.diff(grids[dims[0]])
 
             coords = (midpoints, grids[dims[1]], grids[dims[2]])
@@ -251,17 +251,16 @@ def centroids(self):
         -------
         centroids : numpy.ndarray
             Returns a numpy.ndarray representing the mesh element centroid
-            coordinates with a shape equal to (ndim, dim1, ..., dimn). Can be
-            unpacked along the first dimension with xx, yy, zz = mesh.centroids.
-
-
+            coordinates with a shape equal to (dim1, ..., dimn, ndim). X,
+            Y, Z values can be unpacked with xx, yy, zz =
+            np.rollaxis(mesh.centroids, -1).
         """
         ndim = self.n_dimension
         # this line ensures that the vertices aren't adjusted by the origin or
         # converted to the Cartesian system for cylindrical and spherical meshes
         vertices = StructuredMesh.vertices.fget(self)
-        s0 = (slice(None),) + (slice(0, -1),)*ndim
-        s1 = (slice(None),) + (slice(1, None),)*ndim
+        s0 = (slice(0, -1),)*ndim + (slice(None),)
+        s1 = (slice(1, None),)*ndim + (slice(None),)
         return (vertices[s0] + vertices[s1]) / 2
 
     @property
@@ -351,10 +350,8 @@ def _create_vtk_structured_grid(self):
         import vtk
         from vtk.util import numpy_support as nps
 
-        vertices = self.vertices.T.reshape(-1, 3)
-
         vtkPts = vtk.vtkPoints()
-        vtkPts.SetData(nps.numpy_to_vtk(vertices, deep=True))
+        vtkPts.SetData(nps.numpy_to_vtk(np.swapaxes(self.vertices, 0, 2).reshape(-1, 3), deep=True))
         vtk_grid = vtk.vtkStructuredGrid()
         vtk_grid.SetPoints(vtkPts)
         vtk_grid.SetDimensions(*[dim + 1 for dim in self.dimension])
@@ -373,7 +370,7 @@ def _create_vtk_unstructured_grid(self):
         import vtk
         from vtk.util import numpy_support as nps
 
-        corner_vertices = self.vertices.T.reshape(-1, 3)
+        corner_vertices = np.swapaxes(self.vertices, 0, 2).reshape(-1, 3)
 
         vtkPts = vtk.vtkPoints()
         vtk_grid = vtk.vtkUnstructuredGrid()
@@ -415,7 +412,7 @@ def _insert_point(pnt):
         # list of point IDs
         midpoint_vertices = self.midpoint_vertices
         for edge_grid in midpoint_vertices:
-            for pnt in edge_grid.T.reshape(-1, 3):
+            for pnt in np.swapaxes(edge_grid, 0, 2).reshape(-1, 3):
                 point_ids.append(_insert_point(pnt))
 
         # determine how many elements in each dimension
@@ -448,7 +445,7 @@ def _insert_point(pnt):
                 # initial offset for corner vertices and midpoint dimension
                 flat_idx = corner_vertices.shape[0] + sum(n_midpoint_vertices[:dim])
                 # generate a flat index into the table of point IDs
-                midpoint_shape = midpoint_vertices[dim].shape[1:]
+                midpoint_shape = midpoint_vertices[dim].shape[:-1]
                 flat_idx += np.ravel_multi_index((i+di, j+dj, k+dk),
                                                  midpoint_shape,
                                                  order='F')
@@ -1178,7 +1175,7 @@ class CylindricalMesh(StructuredMesh):
     Parameters
     ----------
     r_grid : numpy.ndarray
-        1-D array of mesh boundary points along the r-axis.
+        1-D array of mesh boundary points along the r-axis
         Requirement is r >= 0.
     z_grid : numpy.ndarray
         1-D array of mesh boundary points along the z-axis relative to the
@@ -1543,14 +1540,14 @@ def centroids_cylindrical(self):
 
     @staticmethod
     def _convert_to_cartesian(arr, origin: Sequence[float]):
-        """Converts an array with xyz values in the first dimension (shape (3, ...))
+        """Converts an array with r, phi, z values in the last dimension (shape (..., 3))
         to Cartesian coordinates.
         """
-        x = arr[0, ...] * np.cos(arr[1, ...]) + origin[0]
-        y = arr[0, ...] * np.sin(arr[1, ...]) + origin[1]
-        arr[0, ...] = x
-        arr[1, ...] = y
-        arr[2, ...] += origin[2]
+        x = arr[..., 0] * np.cos(arr[..., 1]) + origin[0]
+        y = arr[..., 0] * np.sin(arr[..., 1]) + origin[1]
+        arr[..., 0] = x
+        arr[..., 1] = y
+        arr[..., 2] += origin[2]
         return arr
 
 
@@ -1839,16 +1836,16 @@ def centroids_spherical(self):
 
     @staticmethod
     def _convert_to_cartesian(arr, origin: Sequence[float]):
-        """Converts an array with xyz values in the first dimension (shape (3, ...))
+        """Converts an array with r, theta, phi values in the last dimension (shape (..., 3))
         to Cartesian coordinates.
         """
-        r_xy = arr[0, ...] * np.sin(arr[1, ...])
-        x = r_xy * np.cos(arr[2, ...])
-        y = r_xy * np.sin(arr[2, ...])
-        z = arr[0, ...] * np.cos(arr[1, ...])
-        arr[0, ...] = x + origin[0]
-        arr[1, ...] = y + origin[1]
-        arr[2, ...] = z + origin[2]
+        r_xy = arr[..., 0] * np.sin(arr[..., 1])
+        x = r_xy * np.cos(arr[..., 2])
+        y = r_xy * np.sin(arr[..., 2])
+        z = arr[..., 0] * np.cos(arr[..., 1])
+        arr[..., 0] = x + origin[0]
+        arr[..., 1] = y + origin[1]
+        arr[..., 2] = z + origin[2]
         return arr
 
 
@@ -1891,14 +1888,14 @@ class UnstructuredMesh(MeshBase):
     volumes : Iterable of float
         Volumes of the unstructured mesh elements
     centroids : numpy.ndarray
-        Centroids of the mesh elements with array shape (3, n_elements)
+        Centroids of the mesh elements with array shape (n_elements, 3)
 
     vertices : numpy.ndarray
-        Coordinates of the mesh vertices with array shape (3, n_elements)
+        Coordinates of the mesh vertices with array shape (n_elements, 3)
 
         .. versionadded:: 0.13.1
     connectivity : numpy.ndarray
-        Connectivity of the elements with array shape (8, n_elements)
+        Connectivity of the elements with array shape (n_elements, 8)
 
         .. versionadded:: 0.13.1
     element_types : Iterable of integers
@@ -1985,11 +1982,11 @@ def total_volume(self):
 
     @property
     def vertices(self):
-        return self._vertices.T
+        return self._vertices
 
     @property
     def connectivity(self):
-        return self._connectivity.T
+        return self._connectivity
 
     @property
     def element_types(self):
@@ -1997,7 +1994,7 @@ def element_types(self):
 
     @property
     def centroids(self):
-        return np.array([self.centroid(i) for i in range(self.n_elements)]).T
+        return np.array([self.centroid(i) for i in range(self.n_elements)])
 
     @property
     def n_elements(self):
@@ -2053,11 +2050,11 @@ def centroid(self, bin: int):
             x, y, z values of the element centroid
 
         """
-        conn = self.connectivity[:, bin]
+        conn = self.connectivity[bin]
         # remove invalid connectivity values
         conn = conn[conn >= 0]
-        coords = self.vertices[:, conn]
-        return coords.mean(axis=1)
+        coords = self.vertices[conn]
+        return coords.mean(axis=0)
 
     def write_vtk_mesh(self, **kwargs):
         """Map data to unstructured VTK mesh elements.
@@ -2120,12 +2117,12 @@ def write_data_to_vtk(
         grid = vtk.vtkUnstructuredGrid()
 
         vtk_pnts = vtk.vtkPoints()
-        vtk_pnts.SetData(nps.numpy_to_vtk(self.vertices.T))
+        vtk_pnts.SetData(nps.numpy_to_vtk(self.vertices))
         grid.SetPoints(vtk_pnts)
 
         n_skipped = 0
         elems = []
-        for elem_type, conn in zip(self.element_types, self.connectivity.T):
+        for elem_type, conn in zip(self.element_types, self.connectivity):
             if elem_type == self._LINEAR_TET:
                 elem = vtk.vtkTetra()
             elif elem_type == self._LINEAR_HEX:


Test patch:
diff --git a/tests/regression_tests/unstructured_mesh/test.py b/tests/regression_tests/unstructured_mesh/test.py
--- a/tests/regression_tests/unstructured_mesh/test.py
+++ b/tests/regression_tests/unstructured_mesh/test.py
@@ -54,7 +54,7 @@ def _compare_results(self):
                 exp_vertex = (-10.0, -10.0, 10.0)
                 exp_centroid = (-9.0, -9.0, 9.0)
 
-            np.testing.assert_array_equal(umesh.vertices[:, 0], exp_vertex)
+            np.testing.assert_array_equal(umesh.vertices[0], exp_vertex)
             np.testing.assert_array_equal(umesh.centroid(0), exp_centroid)
 
             # loop over the tallies and get data
diff --git a/tests/unit_tests/test_cylindrical_mesh.py b/tests/unit_tests/test_cylindrical_mesh.py
--- a/tests/unit_tests/test_cylindrical_mesh.py
+++ b/tests/unit_tests/test_cylindrical_mesh.py
@@ -98,7 +98,7 @@ def test_offset_mesh(run_in_tmpdir, model, estimator, origin):
         centroids = mesh.centroids
         for ijk in mesh.indices:
             i, j, k = np.array(ijk) - 1
-            if model.geometry.find(centroids[:, i, j, k]):
+            if model.geometry.find(centroids[i, j, k]):
                 mean[i, j, k] == 0.0
             else:
                 mean[i, j, k] != 0.0
diff --git a/tests/unit_tests/test_mesh.py b/tests/unit_tests/test_mesh.py
--- a/tests/unit_tests/test_mesh.py
+++ b/tests/unit_tests/test_mesh.py
@@ -4,6 +4,7 @@
 import pytest
 import openmc
 
+from openmc import RegularMesh, RectilinearMesh, CylindricalMesh, SphericalMesh
 
 @pytest.mark.parametrize("val_left,val_right", [(0, 0), (-1., -1.), (2.0, 2)])
 def test_raises_error_when_flat(val_left, val_right):
@@ -162,29 +163,80 @@ def test_centroids():
     mesh.lower_left = (1., 2., 3.)
     mesh.upper_right = (11., 12., 13.)
     mesh.dimension = (1, 1, 1)
-    np.testing.assert_array_almost_equal(mesh.centroids[:, 0, 0, 0], [6., 7., 8.])
+    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [6., 7., 8.])
 
     # rectilinear mesh
     mesh = openmc.RectilinearMesh()
     mesh.x_grid = [1., 11.]
     mesh.y_grid = [2., 12.]
     mesh.z_grid = [3., 13.]
-    np.testing.assert_array_almost_equal(mesh.centroids[:, 0, 0, 0], [6., 7., 8.])
+    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [6., 7., 8.])
 
     # cylindrical mesh
     mesh = openmc.CylindricalMesh(r_grid=(0, 10), z_grid=(0, 10), phi_grid=(0, np.pi))
-    np.testing.assert_array_almost_equal(mesh.centroids[:, 0, 0, 0], [0.0, 5.0, 5.0])
+    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [0.0, 5.0, 5.0])
     # ensure that setting an origin is handled correctly
     mesh.origin = (5.0, 0, -10)
-    np.testing.assert_array_almost_equal(mesh.centroids[:, 0, 0, 0], [5.0, 5.0, -5.0])
+    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [5.0, 5.0, -5.0])
 
     # spherical mesh, single element xyz-positive octant
     mesh = openmc.SphericalMesh(r_grid=[0, 10], theta_grid=[0, 0.5*np.pi], phi_grid=[0, np.pi])
     x = 5.*np.cos(0.5*np.pi)*np.sin(0.25*np.pi)
     y = 5.*np.sin(0.5*np.pi)*np.sin(0.25*np.pi)
     z = 5.*np.sin(0.25*np.pi)
-    np.testing.assert_array_almost_equal(mesh.centroids[:, 0, 0, 0], [x, y, z])
+    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [x, y, z])
 
     mesh.origin = (-5.0, -5.0, 5.0)
-    np.testing.assert_array_almost_equal(mesh.centroids[:, 0, 0, 0], [x-5.0, y-5.0, z+5.0])
+    np.testing.assert_array_almost_equal(mesh.centroids[0, 0, 0], [x-5.0, y-5.0, z+5.0])
+
+
+@pytest.mark.parametrize('mesh_type', ('regular', 'rectilinear', 'cylindrical', 'spherical'))
+def test_mesh_vertices(mesh_type):
+
+    ijk = (2, 3, 2)
+
+    # create a new mesh object
+    if mesh_type == 'regular':
+        mesh = openmc.RegularMesh()
+        ll = np.asarray([0.]*3)
+        width = np.asarray([0.5]*3)
+        mesh.lower_left = ll
+        mesh.width = width
+        mesh.dimension = (5, 7, 9)
+
+        # spot check that an element has the correct vertex coordinates asociated with it
+        # (using zero-indexing here)
+        exp_i_j_k = ll + np.asarray(ijk, dtype=float) * width
+        np.testing.assert_equal(mesh.vertices[ijk], exp_i_j_k)
+
+        # shift the mesh using the llc
+        shift  = np.asarray((3.0, 6.0, 10.0))
+        mesh.lower_left += shift
+        np.testing.assert_equal(mesh.vertices[ijk], exp_i_j_k+shift)
+    elif mesh_type == 'rectilinear':
+        mesh = openmc.RectilinearMesh()
+        w = np.asarray([0.5] * 3)
+        ll = np.asarray([0.]*3)
+        dims = (5, 7, 9)
+        mesh.x_grid = np.linspace(ll[0], w[0]*dims[0], dims[0])
+        mesh.y_grid = np.linspace(ll[1], w[1]*dims[1], dims[1])
+        mesh.z_grid = np.linspace(ll[2], w[2]*dims[2], dims[2])
+        exp_vert = np.asarray((mesh.x_grid[2], mesh.y_grid[3], mesh.z_grid[2]))
+        np.testing.assert_equal(mesh.vertices[ijk], exp_vert)
+    elif mesh_type == 'cylindrical':
+        r_grid = np.linspace(0, 5, 10)
+        z_grid = np.linspace(-10, 10, 20)
+        phi_grid = np.linspace(0, 2*np.pi, 8)
+        mesh = openmc.CylindricalMesh(r_grid=r_grid, z_grid=z_grid, phi_grid=phi_grid)
+        exp_vert = np.asarray((mesh.r_grid[2], mesh.phi_grid[3], mesh.z_grid[2]))
+        np.testing.assert_equal(mesh.vertices_cylindrical[ijk], exp_vert)
+    elif mesh_type == 'spherical':
+        r_grid = np.linspace(0, 13, 14)
+        theta_grid = np.linspace(0, np.pi, 11)
+        phi_grid = np.linspace(0, 2*np.pi, 7)
+        mesh = openmc.SphericalMesh(r_grid=r_grid, theta_grid=theta_grid, phi_grid=phi_grid)
+        exp_vert = np.asarray((mesh.r_grid[2], mesh.theta_grid[3], mesh.phi_grid[2]))
+        np.testing.assert_equal(mesh.vertices_spherical[ijk], exp_vert)
+
+
 
diff --git a/tests/unit_tests/test_spherical_mesh.py b/tests/unit_tests/test_spherical_mesh.py
--- a/tests/unit_tests/test_spherical_mesh.py
+++ b/tests/unit_tests/test_spherical_mesh.py
@@ -102,7 +102,7 @@ def test_offset_mesh(run_in_tmpdir, model, estimator, origin):
         centroids = mesh.centroids
         for ijk in mesh.indices:
             i, j, k = np.array(ijk) - 1
-            if model.geometry.find(centroids[:, i, j, k]):
+            if model.geometry.find(centroids[i, j, k]):
                 mean[i, j, k] == 0.0
             else:
                 mean[i, j, k] != 0.0