From a93751ce7bcc5589a1dfa1c187c50f9ce6451c42 Mon Sep 17 00:00:00 2001 From: Ethan Peterson Date: Thu, 17 Mar 2022 10:39:22 -0400 Subject: [PATCH] changing meshgrid to vertices --- openmc/mesh.py | 85 ++++++++++++++++++++++++++------------------------ 1 file changed, 44 insertions(+), 41 deletions(-) diff --git a/openmc/mesh.py b/openmc/mesh.py index 76a527357d..056dd34acb 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -245,15 +245,16 @@ class RegularMesh(MeshBase): return ((x,) for x in range(1, nx + 1)) @property - def meshgrid(self): + def vertices(self): """Return coordinates of mesh vertices Returns ------- - coordinates : numpy.ndarray or tuple of numpy.ndarray - Returns a numpy.ndarray for each dimension representing the mesh - vertex coordinates. Each array has an identical shape equal to - (nx + 1, ny + 1, nz + 1) where nx, ny, nz = dimension. + vertices : numpy.ndarray + Returns a numpy.ndarray representing the coordinates of mesh + vertices with a shape equal to (nx + 1, ny + 1, nz + 1, 3) + where nx, ny, nz = dimension in 3D. For 2D the shape is + (nx + 1, ny + 1, 2), and for 1D the shape is (nx,). """ ndim = len(self._dimension) @@ -264,14 +265,16 @@ class RegularMesh(MeshBase): xarr = np.linspace(x0, x1, nx + 1) yarr = np.linspace(y0, y1, ny + 1) zarr = np.linspace(z0, z1, nz + 1) - return np.meshgrid(xarr, yarr, zarr, indexing='ij') + xx, yy, zz = np.meshgrid(xarr, yarr, zarr, indexing='ij') + return np.stack((xx, yy, zz), axis=-1) elif ndim == 2: x0, y0 = self.lower_left x1, y1 = self.upper_right nx, ny = self.dimension xarr = np.linspace(x0, x1, nx + 1) yarr = np.linspace(y0, y1, ny + 1) - return np.meshgrid(xarr, yarr, indexing='ij') + xx, yy = np.meshgrid(xarr, yarr, indexing='ij') + return np.stack((xx, yy), axis=-1) else: nx, = self.dimension x0, = self.lower_left @@ -291,22 +294,13 @@ class RegularMesh(MeshBase): """ ndim = len(self._dimension) - meshgrid = self.meshgrid + vertices = self.vertices if ndim == 3: - xarr, yarr, zarr = meshgrid - xc = (xarr[:-1, :-1, :-1] + xarr[1:, 1:, 1:]) / 2 - yc = (yarr[:-1, :-1, :-1] + yarr[1:, 1:, 1:]) / 2 - zc = (zarr[:-1, :-1, :-1] + zarr[1:, 1:, 1:]) / 2 - return xc, yc, zc + return (vertices[:-1, :-1, :-1, :] + vertices[1:, 1:, 1:, :]) / 2 elif ndim == 2: - xarr, yarr = meshgrid - xc = (xarr[:-1, :-1] + xarr[1:, 1:]) / 2 - yc = (yarr[:-1, :-1] + yarr[1:, 1:]) / 2 - return xc, yc + return (vertices[:-1, :-1, :] + vertices[1:, 1:, :]) / 2 else: - xarr = meshgrid - xc = (xarr[:-1] + xarr[1:]) / 2 - return xc + return (vertices[:-1] + vertices[1:]) / 2 @dimension.setter def dimension(self, dimension): @@ -694,18 +688,19 @@ class RectilinearMesh(MeshBase): for x in range(1, nx + 1)) @property - def meshgrid(self): + def vertices(self): """Return coordinates of mesh vertices Returns ------- - coordinates : 3-tuple of numpy.ndarray - Returns a numpy.ndarray for each dimension representing the mesh - vertex coordinates. Each array has an identical shape equal to - (nx + 1, ny + 1, nz + 1) where nx, ny, nz = dimension. + vertices : numpy.ndarray + Returns a numpy.ndarray representing the coordinates of mesh + vertices with a shape equal to (nx + 1, ny + 1, nz + 1, 3) + where nx, ny, nz = dimension. """ - return np.meshgrid(self.x_grid, self.y_grid, self.z_grid, indexing='ij') + xx, yy, zz = np.meshgrid(self.x_grid, self.y_grid, self.z_grid, indexing='ij') + return np.stack((xx, yy, zz), axis=-1) @property def centroids(self): @@ -719,7 +714,7 @@ class RectilinearMesh(MeshBase): equal to (nx, ny, nz) where nx, ny, nz = n_dimension. """ - xx, yy, zz = self.meshgrid + xx, yy, zz = self.vertices xc = (xx[:-1, :-1, :-1] + xx[1:, 1:, 1:]) / 2 yc = (yy[:-1, :-1, :-1] + yy[1:, 1:, 1:]) / 2 zc = (zz[:-1, :-1, :-1] + zz[1:, 1:, 1:]) / 2 @@ -885,6 +880,10 @@ class CylindricalMesh(MeshBase): def z_grid(self): return self._z_grid + @property + def grids(self): + return (self.r_grid, self.phi_grid, self.z_grid) + @property def indices(self): nr, np, nz = self.dimension @@ -896,18 +895,18 @@ class CylindricalMesh(MeshBase): for r in range(1, nr + 1)) @property - def meshgrid(self): + def vertices(self): """Return coordinates of mesh vertices Returns ------- - coordinates : 3-tuple of numpy.ndarray - Returns a numpy.ndarray for each dimension representing the mesh - vertex coordinates. Each array has an identical shape equal to - (nr + 1, nphi + 1, nz + 1) where nr, nphi, nz = dimension. + vertices : numpy.ndarray + Returns a numpy.ndarray representing the coordinates of mesh + vertices with a shape equal to (nr + 1, nphi + 1, nz + 1, 3) + where nr, nphi, nz = dimension. """ - return np.meshgrid(self.r_grid, self.phi_grid, self.z_grid, indexing='ij') + return np.stack(np.meshgrid(*self.grids, indexing='ij'), axis=-1) @property def centroids(self): @@ -921,7 +920,7 @@ class CylindricalMesh(MeshBase): equal to (nx, nphi, nz) where nx, nphi, nz = n_dimension. """ - rr, pp, zz = self.meshgrid + rr, pp, zz = self.vertices rc = (rr[:-1, :-1, :-1] + rr[1:, 1:, 1:]) / 2 pc = (pp[:-1, :-1, :-1] + pp[1:, 1:, 1:]) / 2 zc = (zz[:-1, :-1, :-1] + zz[1:, 1:, 1:]) / 2 @@ -1116,18 +1115,22 @@ class SphericalMesh(MeshBase): for r in range(1, nr + 1)) @property - def meshgrid(self): + def vertices(self): """Return coordinates of mesh vertices Returns ------- - coordinates : 3-tuple of numpy.ndarray - Returns a numpy.ndarray for each dimension representing the mesh - vertex coordinates. Each array has an identical shape equal to - (nr + 1, ntheta + 1, nphi + 1) where nr, ntheta, nphi = dimension. + vertices : numpy.ndarray + Returns a numpy.ndarray representing the coordinates of mesh + vertices with a shape equal to (nr + 1, ntheta + 1, nphi + 1, 3) + where nr, ntheta, nphi = dimension. """ - return np.meshgrid(self.r_grid, self.phi_grid, self.z_grid, indexing='ij') + rr, tt, pp = np.meshgrid(self.r_grid, + self.theta_grid, + self.phi_grid, + indexing='ij') + return np.stack((rr, tt, pp), axis=-1) @property def centroids(self): @@ -1141,7 +1144,7 @@ class SphericalMesh(MeshBase): equal to (nx, ntheta, nphi) where nx, ntheta, nphi = n_dimension. """ - rr, tt, pp = self.meshgrid + rr, tt, pp = self.vertices rc = (rr[:-1, :-1, :-1] + rr[1:, 1:, 1:]) / 2 tc = (tt[:-1, :-1, :-1] + tt[1:, 1:, 1:]) / 2 pc = (pp[:-1, :-1, :-1] + pp[1:, 1:, 1:]) / 2