From 19f121a45a88f972a0a3fe2c7ec2482101415010 Mon Sep 17 00:00:00 2001 From: Patrick Shriwise Date: Wed, 1 Nov 2023 12:19:27 -0500 Subject: [PATCH] Python mesh coordinates reorder (#2730) Co-authored-by: Ethan Peterson --- openmc/mesh.py | 85 +++++++++---------- .../unstructured_mesh/test.py | 2 +- tests/unit_tests/test_cylindrical_mesh.py | 2 +- tests/unit_tests/test_mesh.py | 64 ++++++++++++-- tests/unit_tests/test_spherical_mesh.py | 2 +- 5 files changed, 102 insertions(+), 53 deletions(-) diff --git a/openmc/mesh.py b/openmc/mesh.py index 4783bd1e1e..a14cec5dc1 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -179,18 +179,18 @@ class StructuredMesh(MeshBase): ------- 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 @@ class StructuredMesh(MeshBase): 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 @@ class StructuredMesh(MeshBase): # 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 @@ class StructuredMesh(MeshBase): ------- 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 @@ class StructuredMesh(MeshBase): 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 @@ class StructuredMesh(MeshBase): 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 @@ class StructuredMesh(MeshBase): # 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 @@ class StructuredMesh(MeshBase): # 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 @@ class CylindricalMesh(StructuredMesh): @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 @@ class SphericalMesh(StructuredMesh): @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 @@ class UnstructuredMesh(MeshBase): @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 @@ class UnstructuredMesh(MeshBase): @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 @@ class UnstructuredMesh(MeshBase): 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 @@ class UnstructuredMesh(MeshBase): 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: diff --git a/tests/regression_tests/unstructured_mesh/test.py b/tests/regression_tests/unstructured_mesh/test.py index 7c7961928e..f3ba46e37a 100644 --- a/tests/regression_tests/unstructured_mesh/test.py +++ b/tests/regression_tests/unstructured_mesh/test.py @@ -54,7 +54,7 @@ class UnstructuredMeshTest(PyAPITestHarness): 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 index df0c5d2bb5..b408bc0e91 100644 --- 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 index 7194519340..06d543610c 100644 --- a/tests/unit_tests/test_mesh.py +++ b/tests/unit_tests/test_mesh.py @@ -4,6 +4,7 @@ import numpy as np 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 index b491014cc7..0b579be35f 100644 --- 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