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)