def plot_image_grid(images: List[np.array],
titles: List[str] = None,
figure: plt.figure = None,
grayscale: bool = False,
transpose: bool = False) -> plt.figure:
"""
Plot a grid of n x m images
Input
-----
images: List[np.array]
Images in a n x m array
titles: List[str] (opt.)
List of m titles for each image column
figure: plt.figure (opt.)
Pyplot figure (if None, will be created)
grayscale: bool (opt.)
If True return grayscaled images
transpose: bool (opt.)
If True, transpose the grid
Output
------
Pyplot figure filled with the images
"""
num_cols, num_rows = len(images), len(images[0])
img_ratio = images[0][0].shape[1] / images[0][0].shape[0]
if transpose:
vert_grid_shape, hori_grid_shape = (1, num_rows), (num_cols, 1)
figsize = (int(num_rows * 5 * img_ratio), num_cols * 5)
wspace, hspace = 0.2, 0.
else:
vert_grid_shape, hori_grid_shape = (num_rows, 1), (1, num_cols)
figsize = (int(num_cols * 5 * img_ratio), num_rows * 5)
hspace, wspace = 0.2, 0.
if figure is None:
figure = plt.figure(figsize=figsize)
imshow_params = {'cmap': plt.get_cmap('gray')} if grayscale else {}
grid_spec = gridspec.GridSpec(*hori_grid_shape, wspace=0, hspace=0)
for j in range(num_cols):
grid_spec_j = gridspec.GridSpecFromSubplotSpec(
*vert_grid_shape,
subplot_spec=grid_spec[j],
wspace=wspace,
hspace=hspace)
for i in range(num_rows):
ax_img = figure.add_subplot(grid_spec_j[i])
# ax_img.axis('off')
ax_img.set_yticks([])
ax_img.set_xticks([])
if titles is not None:
if transpose:
ax_img.set_ylabel(titles[j], fontsize=25)
else:
ax_img.set_title(titles[j], fontsize=15)
ax_img.imshow(images[j][i], **imshow_params)
figure.tight_layout()
return figure
def Update_Axes(fig: plt.figure, Axes: np.ndarray, u_NN: Neural_Network,
Points: torch.Tensor, n: int) -> None:
""" This function plots the approximate solution and residual at the
specified points.
----------------------------------------------------------------------------
Arguments:
fig : The figure object to which the Axes belong. We need this to set up
the color bars.
Axes : The array of Axes object that we will plot on. Note that this
function will overwrite these axes.
u_NN : The neural network that gives the approximate solution to Poisson's
equation.
Points : The set of points we want to evaluate the approximate and true
solutions, as well as the PDE Residual. This should be an (n*n)x2 tensor,
whose ith row holds the x,y coordinates of the ith point we want to plot.
Each element of Points should be an element of [0,1]x[0,1].
n : the number of gridpoints along each axis. Points should be an n*n x 2
tensor.
----------------------------------------------------------------------------
Returns:
Nothing! """
# First, evaluate the network's approximate solution, the true solution, and
# the PDE residual at the specified Points. We need to reshape these into
# nxn grids, because that's what matplotlib's contour function wants. It's
# annoying, but it is what it is.
u_NN_at_Points = Evaluate_NN(u_NN, Points).reshape(n, n)
True_Sol_at_Points = Evaluate_True_Solution(Points).reshape(n, n)
Residual_at_Points = PDE_Residual(u_NN, Points).reshape(n, n)
# Extract the x and y coordinates of points, as np arrays. We also need to
# reshape these as nxn grids (same reason as above.
x = Points[:, 0].numpy().reshape(n, n)
y = Points[:, 1].numpy().reshape(n, n)
# Plot the approximate solution + colorbar.
ColorMap0 = Axes[0].contourf(x,
y,
u_NN_at_Points.reshape(n, n),
levels=50,
cmap=plt.cm.jet)
fig.colorbar(ColorMap0,
ax=Axes[0],
fraction=0.046,
pad=0.04,
orientation='vertical')
# Plot the true solution + colorbar
ColorMap1 = Axes[1].contourf(x,
y,
True_Sol_at_Points.reshape(n, n),
levels=50,
cmap=plt.cm.jet)
fig.colorbar(ColorMap1,
ax=Axes[1],
fraction=0.046,
pad=0.04,
orientation='vertical')
# Plot the residual + colorbar
ColorMap2 = Axes[2].contourf(x,
y,
Residual_at_Points.reshape(n, n),
levels=50,
cmap=plt.cm.jet)
fig.colorbar(ColorMap2,
ax=Axes[2],
fraction=0.046,
pad=0.04,
orientation='vertical')
# Set tight layout (to prevent overlapping... I have no idea why this isn't
# a default setting. Matplotlib, you are special kind of awful).
fig.tight_layout()
def Plot_Solution(
fig : plt.figure,
Axes : numpy.ndarray,
Sol_NN : Neural_Network,
PDE_NN : Neural_Network,
Time_Derivative_Order : int,
Spatial_Derivative_Order : int,
Data : Data_Container) -> None:
""" This function makes four plots. One for the approximate solution, one
for the true solution, one for their difference, and one for the PDE
residual. x_points and t_points specify the domain of all four plots.
Note: this function only works if u is a function of 1 spatial variable.
----------------------------------------------------------------------------
Arguments:
fig: The figure object to which the Axes belong. We need this to set up
the color bars.
Axes: The array of Axes object that we will plot on. Note that this
function overwrites these axes.
Sol_NN: The network that approximates the PDE solution.
PDE_NN: The network that approximates the PDE.
Time_Derivative_Order: The order of the time derivative on the left-hand
side of the PDE.
Spatial_Derivative_Order: The highest order spatial derivatives of Sol_NN we
need to evaluate.
Data: This is a Data_Container object. It should contain four members (all
of which are numpy arrays): x_points, t_points, Data_Set, and Noisy_Data_Set.
x_points, t_points contain the set of possible x and t values, respectively.
Data_Set and Noisy_Data_Set should contain the true solution with and
without noise at each grid point (t, x coordinate). If t_points and x_points
have n_t and n_x elements, respectively, then Data_Set and Noisy_Data_Set
should be an n_x by n_t array whose i,j element holds the value of the true
solution at t_points[j], x_points[i].
----------------------------------------------------------------------------
Returns:
Nothing! """
# First, construct the set of possible coordinates. grid_t_coords and
# grid_x_coords are 2d NumPy arrays. Each row of these arrays corresponds to
# a specific position. Each column corresponds to a specific time.
grid_t_coords, grid_x_coords = numpy.meshgrid(Data.t_points, Data.x_points);
# Flatten t_coords, x_coords. use them to generate grid point coodinates.
flattened_grid_x_coords = grid_x_coords.reshape(-1, 1);
flattened_grid_t_coords = grid_t_coords.reshape(-1, 1);
Grid_Point_Coords = torch.from_numpy(numpy.hstack((flattened_grid_t_coords, flattened_grid_x_coords)));
# Get number of x and t values, respectively.
n_x = len(Data.x_points);
n_t = len(Data.t_points);
# Put networks into evaluation mode.
Sol_NN.eval();
PDE_NN.eval();
# Evaluate the network's approximate solution, the absolute error, and the
# PDE residual at each coordinate. We need to reshape these into n_x by n_t
# grids because that's what matplotlib's contour function wants.
Approx_Sol_on_grid = Evaluate_Approx_Sol(Sol_NN, Grid_Point_Coords).reshape(n_x, n_t);
Error_On_Grid = numpy.abs(Approx_Sol_on_grid - Data.Data_Set);
Residual_on_Grid = Evaluate_Residual(
Sol_NN = Sol_NN,
PDE_NN = PDE_NN,
Time_Derivative_Order = Time_Derivative_Order,
Spatial_Derivative_Order = Spatial_Derivative_Order,
Coords = Grid_Point_Coords).reshape(n_x, n_t);
# Plot the true solution + color bar.
data_min : float = numpy.min(Data.Noisy_Data_Set);
data_max : float = numpy.max(Data.Noisy_Data_Set);
ColorMap0 = Axes[0].contourf( grid_t_coords,
grid_x_coords,
Data.Noisy_Data_Set,
levels = numpy.linspace(data_min, data_max, 500),
cmap = plt.cm.jet);
fig.colorbar(ColorMap0, ax = Axes[0], fraction=0.046, pad=0.04, orientation='vertical');
# Plot the learned solution + color bar
sol_min : float = numpy.min(Approx_Sol_on_grid);
sol_max : float = numpy.max(Approx_Sol_on_grid);
ColorMap1 = Axes[1].contourf( grid_t_coords,
grid_x_coords,
Approx_Sol_on_grid,
levels = numpy.linspace(sol_min, sol_max, 500),
cmap = plt.cm.jet);
fig.colorbar(ColorMap1, ax = Axes[1], fraction=0.046, pad=0.04, orientation='vertical');
# Plot the Error between the approx solution and noise-free data set. + color bar.
error_min : float = numpy.min(Error_On_Grid);
error_max : float = numpy.max(Error_On_Grid);
ColorMap2 = Axes[2].contourf( grid_t_coords,
grid_x_coords,
Error_On_Grid,
levels = numpy.linspace(error_min, error_max, 500),
cmap = plt.cm.jet);
fig.colorbar(ColorMap2, ax = Axes[2], fraction=0.046, pad=0.04, orientation='vertical');
# Plot the residual + color bar
resid_min : float = numpy.min(Residual_on_Grid);
resid_max : float = numpy.max(Residual_on_Grid);
ColorMap3 = Axes[3].contourf( grid_t_coords,
grid_x_coords,
Residual_on_Grid,
levels = numpy.linspace(resid_min, resid_max, 500),
cmap = plt.cm.jet);
fig.colorbar(ColorMap3, ax = Axes[3], fraction=0.046, pad=0.04, orientation='vertical');
# Set tight layout (to prevent overlapping... I have no idea why this isn't
# a default setting. Matplotlib, you are special kind of awful).
fig.tight_layout();
