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)