def test_spherical_mesh():
mesh = Mesh([(1., 2.), (0., 2 * np.pi), (0., np.pi)],
[.5, np.pi, np.pi / 2.], CoordinateSystem.SPHERICAL)
expected_spherical_vertex_coordinate_grids = [
np.array([[[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]],
[[1.5, 1.5, 1.5], [1.5, 1.5, 1.5], [1.5, 1.5, 1.5]],
[[2., 2., 2.], [2., 2., 2.], [2., 2., 2.]]]),
np.array([[[0., 0., 0.], [np.pi, np.pi, np.pi],
[2. * np.pi, 2. * np.pi, 2. * np.pi]],
[[0., 0., 0.], [np.pi, np.pi, np.pi],
[2. * np.pi, 2. * np.pi, 2. * np.pi]],
[[0., 0., 0.], [np.pi, np.pi, np.pi],
[2. * np.pi, 2. * np.pi, 2. * np.pi]]]),
np.array([[[0., np.pi / 2., np.pi], [0., np.pi / 2., np.pi],
[0., np.pi / 2., np.pi]],
[[0., np.pi / 2., np.pi], [0., np.pi / 2., np.pi],
[0., np.pi / 2., np.pi]],
[[0., np.pi / 2., np.pi], [0., np.pi / 2., np.pi],
[0., np.pi / 2., np.pi]]])
]
actual_spherical_vertex_coordinate_grids = mesh.coordinate_grids(True)
assert np.allclose(actual_spherical_vertex_coordinate_grids,
expected_spherical_vertex_coordinate_grids)
expected_spherical_cell_center_coordinate_grids = [
np.array([[[1.25, 1.25], [1.25, 1.25]], [[1.75, 1.75], [1.75, 1.75]]]),
np.array([[[np.pi / 2., np.pi / 2.],
[3. * np.pi / 2., 3. * np.pi / 2.]],
[[np.pi / 2., np.pi / 2.],
[3. * np.pi / 2., 3. * np.pi / 2.]]]),
np.array([[[np.pi / 4., 3. * np.pi / 4.],
[np.pi / 4., 3. * np.pi / 4.]],
[[np.pi / 4., 3. * np.pi / 4.],
[np.pi / 4., 3. * np.pi / 4.]]])
]
actual_spherical_cell_center_coordinate_grids = \
mesh.coordinate_grids(False)
assert np.allclose(actual_spherical_cell_center_coordinate_grids,
expected_spherical_cell_center_coordinate_grids)
expected_cartesian_vertex_coordinate_grids = [
np.array([[[0., 1., 0.], [0., -1., 0.], [0., 1., 0.]],
[[0., 1.5, 0.], [0., -1.5, 0.], [0., 1.5, 0.]],
[[0., 2., 0.], [0., -2., 0.], [0., 2., 0.]]]),
np.array([[[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]],
[[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]],
[[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]]]),
np.array([[[1., 0., -1.], [1., 0., -1.], [1., 0., -1.]],
[[1.5, 0., -1.5], [1.5, 0., -1.5], [1.5, 0., -1.5]],
[[2., 0., -2.], [2., 0., -2.], [2., 0., -2.]]])
]
actual_cartesian_vertex_coordinate_grids = \
mesh.cartesian_coordinate_grids(True)
assert np.allclose(actual_cartesian_vertex_coordinate_grids,
expected_cartesian_vertex_coordinate_grids)
expected_cartesian_cell_center_coordinate_grids = [
np.array([[[0., 0.], [0., 0.]], [[0., 0.], [0., 0.]]]),
np.array([[[.88388348, .88388348], [-.88388348, -.88388348]],
[[1.23743687, 1.23743687], [-1.23743687, -1.23743687]]]),
np.array([[[.88388348, -.88388348], [.88388348, -.88388348]],
[[1.23743687, -1.23743687], [1.23743687, -1.23743687]]]),
]
actual_cartesian_cell_center_coordinate_grids = \
mesh.cartesian_coordinate_grids(False)
assert np.allclose(actual_cartesian_cell_center_coordinate_grids,
expected_cartesian_cell_center_coordinate_grids)
expected_vertex_unit_vector_grids = [
np.array([[[[0., 0., 1.], [1., 0., 0.], [0., 0., -1.]],
[[0., 0., 1.], [-1., 0., 0.], [0., 0., -1.]],
[[0., 0., 1.], [1., 0., 0.], [0., 0., -1.]]],
[[[0., 0., 1.], [1., 0., 0.], [0., 0., -1.]],
[[0., 0., 1.], [-1., 0., 0.], [0., 0., -1.]],
[[0., 0., 1.], [1., 0., 0.], [0., 0., -1.]]],
[[[0., 0., 1.], [1., 0., 0.], [0., 0., -1.]],
[[0., 0., 1.], [-1., 0., 0.], [0., 0., -1.]],
[[0., 0., 1.], [1., 0., 0.], [0., 0., -1.]]]]),
np.array([[[[0., 1., 0.], [0., 1., 0.], [0., 1., 0.]],
[[0., -1., 0.], [0., -1., 0.], [0., -1., 0.]],
[[0., 1., 0.], [0., 1., 0.], [0., 1., 0.]]],
[[[0., 1., 0.], [0., 1., 0.], [0., 1., 0.]],
[[0., -1., 0.], [0., -1., 0.], [0., -1., 0.]],
[[0., 1., 0.], [0., 1., 0.], [0., 1., 0.]]],
[[[0., 1., 0.], [0., 1., 0.], [0., 1., 0.]],
[[0., -1., 0.], [0., -1., 0.], [0., -1., 0.]],
[[0., 1., 0.], [0., 1., 0.], [0., 1., 0.]]]]),
np.array([[[[1., 0., 0.], [0., 0., -1.], [-1., 0., 0.]],
[[-1., 0., 0.], [0., 0., -1.], [1., 0., 0.]],
[[1., 0., 0.], [0., 0., -1.], [-1., 0., 0.]]],
[[[1., 0., 0.], [0., 0., -1.], [-1., 0., 0.]],
[[-1., 0., 0.], [0., 0., -1.], [1., 0., 0.]],
[[1., 0., 0.], [0., 0., -1.], [-1., 0., 0.]]],
[[[1., 0., 0.], [0., 0., -1.], [-1., 0., 0.]],
[[-1., 0., 0.], [0., 0., -1.], [1., 0., 0.]],
[[1., 0., 0.], [0., 0., -1.], [-1., 0., 0.]]]])
]
actual_vertex_unit_vector_grids = mesh.unit_vector_grids(True)
assert np.allclose(actual_vertex_unit_vector_grids,
expected_vertex_unit_vector_grids)
expected_cell_center_unit_vector_grids = [
np.array([[[[0., .707106781, .707106781],
[0., .707106781, -.707106781]],
[[0., -.707106781, .707106781],
[0., -.707106781, -.707106781]]],
[[[0., .707106781, .707106781],
[0., .707106781, -.707106781]],
[[0., -.707106781, .707106781],
[0., -.707106781, -.707106781]]]]),
np.array([[[[-1., 0., 0.], [-1., 0., 0.]], [[1., 0., 0.], [1., 0.,
0.]]],
[[[-1., 0., 0.], [-1., 0., 0.]], [[1., 0., 0.], [1., 0.,
0.]]]]),
np.array([[[[0., .707106781, -.707106781],
[0., -.707106781, -.707106781]],
[[0., -.707106781, -.707106781],
[0., .707106781, -.707106781]]],
[[[0., .707106781, -.707106781],
[0., -.707106781, -.707106781]],
[[0., -.707106781, -.707106781],
[0., .707106781, -.707106781]]]])
]
actual_cell_center_unit_vector_grids = mesh.unit_vector_grids(False)
assert np.allclose(actual_cell_center_unit_vector_grids,
expected_cell_center_unit_vector_grids)
def test_polar_mesh():
mesh = Mesh([(0., 1.), (0., np.pi)], [.5, np.pi / 2.],
CoordinateSystem.POLAR)
expected_polar_vertex_coordinate_grids = [
np.array([
[0., 0., 0.],
[.5, .5, .5],
[1., 1., 1.],
]),
np.array([
[0., np.pi / 2., np.pi],
[0., np.pi / 2., np.pi],
[0., np.pi / 2., np.pi],
])
]
actual_polar_vertex_coordinate_grids = mesh.coordinate_grids(True)
assert np.allclose(actual_polar_vertex_coordinate_grids,
expected_polar_vertex_coordinate_grids)
expected_polar_cell_center_coordinate_grids = [
np.array([
[.25, .25],
[.75, .75],
]),
np.array([
[np.pi / 4., 3. * np.pi / 4.],
[np.pi / 4., 3. * np.pi / 4.],
])
]
actual_polar_cell_center_coordinate_grids = mesh.coordinate_grids(False)
assert np.allclose(actual_polar_cell_center_coordinate_grids,
expected_polar_cell_center_coordinate_grids)
expected_cartesian_vertex_coordinate_grids = [
np.array([
[0., 0., 0.],
[.5, 0., -.5],
[1., 0., -1.],
]),
np.array([
[0., 0., 0.],
[0., .5, 0.],
[0., 1., 0.],
])
]
actual_cartesian_vertex_coordinate_grids = \
mesh.cartesian_coordinate_grids(True)
assert np.allclose(actual_cartesian_vertex_coordinate_grids,
expected_cartesian_vertex_coordinate_grids)
expected_cartesian_cell_center_coordinate_grids = [
np.array([[.1767767, -.1767767], [.53033009, -.53033009]]),
np.array([[.1767767, .1767767], [.53033009, .53033009]])
]
actual_cartesian_cell_center_coordinate_grids = \
mesh.cartesian_coordinate_grids(False)
assert np.allclose(actual_cartesian_cell_center_coordinate_grids,
expected_cartesian_cell_center_coordinate_grids)
expected_vertex_unit_vector_grids = [
np.array([[[1., 0.], [0., 1.], [-1., 0.]],
[[1., 0.], [0., 1.], [-1., 0.]],
[[1., 0.], [0., 1.], [-1., 0.]]]),
np.array([[[0., 1.], [-1., 0.], [0., -1.]],
[[0., 1.], [-1., 0.], [0., -1.]],
[[0., 1.], [-1., 0.], [0., -1.]]])
]
actual_vertex_unit_vector_grids = mesh.unit_vector_grids(True)
assert np.allclose(actual_vertex_unit_vector_grids,
expected_vertex_unit_vector_grids)
expected_cell_center_unit_vector_grids = [
np.array([[[.70710678, .70710678], [-.70710678, .70710678]],
[[.70710678, .70710678], [-.70710678, .70710678]]]),
np.array([[[-.70710678, .70710678], [-.70710678, -.70710678]],
[[-.70710678, .70710678], [-.70710678, -.70710678]]])
]
actual_cell_center_unit_vector_grids = mesh.unit_vector_grids(False)
assert np.allclose(actual_cell_center_unit_vector_grids,
expected_cell_center_unit_vector_grids)
def __init__(
self,
y: np.ndarray,
mesh: Mesh,
vertex_oriented: bool,
n_frames: int = 100,
interval: int = 100,
normalize: bool = False,
pivot: str = 'middle',
quiver_scale: float = 10.,
**_):
"""
:param y: an array representing the solution vector field of the
partial differential equation system
:param mesh: the spatial mesh over which the solution is evaluated
:param vertex_oriented: whether the solution is evaluated over the
vertices or the cell centers of the mesh
:param n_frames: the number of frames to display
:param interval: the number of milliseconds to pause between each frame
:param normalize: Wheter to normalize the lengths of the arrows to one
:param pivot: the pivot point of the arrows
:param quiver_scale: the scaling factor to apply to the arrow lengths
:param _: any ignored extra arguments
"""
self._verify_pde_solution_shape_matches_problem(
y, mesh, vertex_oriented, (2, 3), True)
x_cartesian_coordinate_grids = mesh.cartesian_coordinate_grids(
vertex_oriented)
unit_vector_grids = mesh.unit_vector_grids(vertex_oriented)
y_cartesian: np.ndarray = np.asarray(sum([
y[..., i:i + 1] * unit_vector_grids[i][np.newaxis, ...]
for i in range(mesh.dimensions)
]))
self._quiver_plot: Optional[Quiver] = None
fig = plt.figure()
if mesh.dimensions == 2:
y_0 = y_cartesian[..., 0]
y_1 = y_cartesian[..., 1]
if normalize:
y_magnitude = np.sqrt(np.square(y_0) + np.square(y_1))
y_magnitude_gt_zero = y_magnitude > 0.
y_0[y_magnitude_gt_zero] /= y_magnitude[y_magnitude_gt_zero]
y_1[y_magnitude_gt_zero] /= y_magnitude[y_magnitude_gt_zero]
ax = fig.add_subplot()
def init_plot():
ax.clear()
ax.set_xlabel('x')
ax.set_ylabel('y')
self._quiver_plot = ax.quiver(
*x_cartesian_coordinate_grids,
y_0[0, ...],
y_1[0, ...],
pivot=pivot,
angles='xy',
scale_units='xy',
scale=1. / quiver_scale)
ax.axis('scaled')
def update_plot(time_step: int):
self._quiver_plot.set_UVC(
y_0[time_step, ...],
y_1[time_step, ...])
else:
y_0 = y_cartesian[..., 0] * quiver_scale
y_1 = y_cartesian[..., 1] * quiver_scale
y_2 = y_cartesian[..., 2] * quiver_scale
ax = fig.add_subplot(projection='3d')
def init_plot():
ax.clear()
self._quiver_plot = ax.quiver(
*x_cartesian_coordinate_grids,
y_0[0, ...],
y_1[0, ...],
y_2[0, ...],
pivot=pivot,
normalize=normalize)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
ax.set_box_aspect((
np.ptp(x_cartesian_coordinate_grids[0]),
np.ptp(x_cartesian_coordinate_grids[1]),
np.ptp(x_cartesian_coordinate_grids[2])))
def update_plot(time_step: int):
self._quiver_plot.remove()
self._quiver_plot = ax.quiver(
*x_cartesian_coordinate_grids,
y_0[time_step, ...],
y_1[time_step, ...],
y_2[time_step, ...],
pivot=pivot,
normalize=normalize)
super(QuiverPlot, self).__init__(
fig, init_plot, update_plot, y.shape[0], n_frames, interval)
