Python mesh coordinates reorder (#2730)

Co-authored-by: Ethan Peterson <eepeterson3@gmail.com>
This commit is contained in:
Patrick Shriwise 2023-11-01 12:19:27 -05:00 committed by GitHub
parent ac941f79e0
commit 19f121a45a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 102 additions and 53 deletions

View file

@ -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:

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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