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 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 __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)
def __init__(
self,
y: np.ndarray,
mesh: Mesh,
vertex_oriented: bool,
n_frames: int = 100,
interval: int = 100,
color_map: Colormap = cm.viridis,
v_min: Optional[float] = None,
v_max: Optional[float] = None,
marker_shape: str = 'o',
marker_size: Union[float, np.ndarray] = 20.,
marker_opacity: float = 1.,
**_):
"""
:param y: an array representing the solution of the 3D partial
differential equation
: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_map: the color map to use to map the values of the
solution scalar field to colors
:param v_min: the lower limit of the color map; if None, the limit is
set to the minimum of the solution
:param v_max: the upper limit of the color map; if None, the limit is
set to the maximum of the solution
:param marker_shape: the shape of the point markers
:param marker_size: the size of the point markers
:param marker_opacity: the opacity of the point markers
:param _: any ignored extra arguments
"""
self._verify_pde_solution_shape_matches_problem(
y, mesh, vertex_oriented, 3, False)
x_cartesian_coordinate_grids = \
mesh.cartesian_coordinate_grids(vertex_oriented)
mappable = ScalarMappable(cmap=color_map)
mappable.set_clim(
np.min(y) if v_min is None else v_min,
np.max(y) if v_max is None else v_max)
self._scatter_plot: Optional[PathCollection] = None
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
def init_plot():
ax.clear()
ax.set_xlabel('x0')
ax.set_ylabel('x1')
ax.set_zlabel('x2')
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])))
self._scatter_plot = ax.scatter(
*x_cartesian_coordinate_grids,
c=mappable.to_rgba(y[0, ..., 0].flatten()),
marker=marker_shape,
s=marker_size,
alpha=marker_opacity)
def update_plot(time_step: int):
self._scatter_plot.set_color(
mappable.to_rgba(y[time_step, ..., 0].flatten()))
super(ScatterPlot, self).__init__(
fig, init_plot, update_plot, y.shape[0], n_frames, interval)
def __init__(
self,
y: np.ndarray,
mesh: Mesh,
vertex_oriented: bool,
n_frames: int = 100,
interval: int = 100,
color_map: Colormap = cm.viridis,
v_min: Optional[float] = None,
v_max: Optional[float] = None,
equal_scale: bool = False,
**_):
"""
:param y: an array representing the solution scalar field of the 2D
partial differential equation
: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_map: the color map to use to map the values of the
solution scalar field to colors
:param v_min: the lower z-axis and color map limit; if None, both of
these limits are set to the minimum of the solution
:param v_max: the upper z-axis and color map limit; if None, both of
these limits are set to the maximum of the solution
:param equal_scale: whether the scale of the values of the solution
scalar field is the same as the scale of the spatial dimensions
(i.e. the values represent height)
:param _: any ignored extra arguments
"""
self._verify_pde_solution_shape_matches_problem(
y, mesh, vertex_oriented, 2, False)
x_cartesian_coordinate_grids = \
mesh.cartesian_coordinate_grids(vertex_oriented)
v_min = np.min(y) if v_min is None else v_min
v_max = np.max(y) if v_max is None else v_max
x_0_ptp = np.ptp(x_cartesian_coordinate_grids[0])
x_1_ptp = np.ptp(x_cartesian_coordinate_grids[1])
x_2_ptp = (v_max - v_min) if equal_scale else min(x_0_ptp, x_1_ptp)
surface_plot_args = {
'vmin': v_min,
'vmax': v_max,
'rstride': 1,
'cstride': 1,
'linewidth': 0,
'antialiased': False,
'cmap': color_map
}
self._surface_plot: Optional[Poly3DCollection] = None
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
def init_plot():
ax.clear()
self._surface_plot = ax.plot_surface(
*x_cartesian_coordinate_grids,
y[0, ..., 0],
**surface_plot_args)
ax.set_xlabel('x0')
ax.set_ylabel('x1')
ax.set_zlabel('y')
ax.set_zlim(v_min, v_max)
ax.set_box_aspect((x_0_ptp, x_1_ptp, x_2_ptp))
def update_plot(time_step: int):
self._surface_plot.remove()
self._surface_plot = ax.plot_surface(
*x_cartesian_coordinate_grids,
y[time_step, ..., 0],
**surface_plot_args)
super(SurfacePlot, self).__init__(
fig, init_plot, update_plot, y.shape[0], n_frames, interval)
def __init__(
self,
y: np.ndarray,
mesh: Mesh,
vertex_oriented: bool,
n_frames: int = 100,
interval: int = 100,
color_map: Colormap = cm.viridis,
v_min: Optional[float] = None,
v_max: Optional[float] = None,
**_):
"""
:param y: an array representing the solution scalar field of the 2D
partial differential equation
: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_map: the color map to use to map the values of the
solution scalar field to colors
:param v_min: the lower limit of the color map; if None, the limit is
set to the minimum of the solution
:param v_max: the upper limit of the color map; if None, the limit is
set to the maximum of the solution
:param _: any ignored extra arguments
"""
self._verify_pde_solution_shape_matches_problem(
y, mesh, vertex_oriented, 2, False)
x_cartesian_coordinate_grids = \
mesh.cartesian_coordinate_grids(vertex_oriented)
v_min = np.min(y) if v_min is None else v_min
v_max = np.max(y) if v_max is None else v_max
self._contour_plot: Optional[ContourSet] = None
fig = plt.figure()
def init_plot():
fig.clear()
ax = fig.add_subplot()
self._contour_plot = ax.contourf(
*x_cartesian_coordinate_grids,
y[0, ..., 0],
vmin=v_min,
vmax=v_max,
cmap=color_map)
ax.set_xlabel('x0')
ax.set_ylabel('x1')
ax.axis('scaled')
mappable = ScalarMappable(cmap=color_map)
mappable.set_clim(v_min, v_max)
plt.colorbar(mappable=mappable)
def update_plot(time_step: int):
for collection in self._contour_plot.collections:
collection.remove()
self._contour_plot = self._contour_plot.axes.contourf(
*x_cartesian_coordinate_grids,
y[time_step, ..., 0],
vmin=v_min,
vmax=v_max,
cmap=color_map)
super(ContourPlot, self).__init__(
fig, init_plot, update_plot, y.shape[0], n_frames, interval)
