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,
color: str = 'black',
density: float = 1.,
**_):
"""
:param y: an array representing the solution vector field of the 2D
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 color: the color to use for the lines and arrows of the stream
plot
:param density: the density of the stream lines
:param _: any ignored extra arguments
"""
self._verify_pde_solution_shape_matches_problem(
y, mesh, vertex_oriented, 2, True)
coordinate_grids = mesh.coordinate_grids(vertex_oriented)
self._stream_plot: Optional[StreamplotSet] = None
fig = plt.figure()
if mesh.coordinate_system_type == CoordinateSystem.POLAR:
(x_1_min, x_1_max), (x_0_min, x_0_max) = mesh.x_intervals
x_1_min = 0
x_0 = coordinate_grids[1]
x_1 = coordinate_grids[0]
y_0 = y[..., 1]
y_1 = y[..., 0]
ax = fig.add_subplot(projection='polar')
else:
(x_0_min, x_0_max), (x_1_min, x_1_max) = mesh.x_intervals
x_0 = coordinate_grids[0].T
x_1 = coordinate_grids[1].T
y_0 = y[..., 0].transpose([0, 2, 1])
y_1 = y[..., 1].transpose([0, 2, 1])
ax = fig.add_subplot()
def init_plot():
ax.clear()
self._stream_plot = ax.streamplot(
x_0,
x_1,
y_0[0, ...],
y_1[0, ...],
color=color,
density=density)
ax.set_xlim(x_0_min, x_0_max)
ax.set_ylim(x_1_min, x_1_max)
if mesh.coordinate_system_type == CoordinateSystem.CARTESIAN:
ax.axis('scaled')
ax.set_xlabel('x')
ax.set_ylabel('y')
def update_plot(time_step: int):
with warnings.catch_warnings():
warnings.simplefilter('ignore')
ax.patches.clear()
self._stream_plot.lines.remove()
self._stream_plot = ax.streamplot(
x_0,
x_1,
y_0[time_step, ...],
y_1[time_step, ...],
color=color,
density=density)
super(StreamPlot, self).__init__(
fig, init_plot, update_plot, y.shape[0], n_frames, interval)
