def histplot_2d(
var_x: pd.Series, var_y: pd.Series,
xbins: Union[tuple, list], ybins: Union[tuple, list],
weights: pd.Series,
ax: plt.axes,
fig: plt.figure,
n_threads: int = config.n_threads,
is_z_log: bool = True,
is_square: bool = True,
) -> bh.Histogram:
"""
Plots and prints out 2d histogram. Does not support axis transforms (yet!)
:param var_x: pandas series of var to plot on x-axis
:param var_y: pandas series of var to plot on y-axis
:param xbins: tuple of bins in x (n_bins, start, stop) or list of bin edges
:param ybins: tuple of bins in y (n_bins, start, stop) or list of bin edges
:param weights: series of weights to apply to axes
:param ax: axis to plot on
:param fig: figure to plot on (for colourbar)
:param n_threads: number of threads for filling
:param is_z_log: whether z-axis should be scaled logarithmically
:param is_square: whether to set square aspect ratio
:return: histogram
"""
# setup and fill histogram
hist_2d = bh.Histogram(get_axis(xbins), get_axis(ybins))
hist_2d.fill(var_x, var_y, weight=weights, threads=n_threads)
if is_z_log:
norm = LogNorm()
else:
norm = None
# setup mesh differently depending on bh storage
if hasattr(hist_2d.view(), 'value'):
mesh = ax.pcolormesh(*hist_2d.axes.edges.T, hist_2d.view().value.T, norm=norm)
else:
mesh = ax.pcolormesh(*hist_2d.axes.edges.T, hist_2d.view().T, norm=norm)
fig.colorbar(mesh, ax=ax, fraction=0.046, pad=0.04)
if is_square: # square aspect ratio
ax.set_aspect(1 / ax.get_data_ratio())
return hist_2d
def draw_model_prediction(fig: plt.figure,
ax: plt.axis,
model: nn.Module,
gw: int = 0,
min_point: (float, float) = (0, 0),
max_point: (float, float) = (500, 500),
resolution: float = 1,
apply_ploss: bool = True,
colorbar: bool = True,
detection_flag: bool = False,
detection_threshold: float = -140,
device='cpu',
*args,
**kwargs):
x_mesh, y_mesh = np.mgrid[min_point[0]:max_point[0]:resolution,
min_point[1]:max_point[1]:resolution]
x = x_mesh.ravel()
y = y_mesh.ravel()
data_x = np.stack([x, y], axis=1).astype(np.float32)
if detection_flag is False:
z = model(torch.from_numpy(data_x).to(device),
apply_ploss).detach()[:, gw, 0].cpu().numpy()
else:
#z = model(torch.from_numpy(data_x).to(device), apply_ploss).detach()[:, gw, 0].cpu().numpy()
z = model(torch.from_numpy(data_x).to(device),
apply_ploss).detach()[:, gw, :].cpu().numpy()
sub_z1 = z[:, 1]
sub_z = z[:, 0]
sub_z[sub_z1 < 0.9] = detection_threshold
z = z[:, 0]
z = z.reshape((int(max_point[0] - min_point[0]),
int(max_point[1] - min_point[1]))).transpose()
z[z < detection_threshold] = detection_threshold
im = ax.imshow(z, norm=colors.PowerNorm(gamma=2))
if colorbar:
divider = make_axes_locatable(ax)
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(im, cax=cax, orientation='vertical', label='dBm')
return im
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();
