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)