import torch import matplotlib from matplotlib.colors import ListedColormap import seaborn def get_random_messages(count, order): """ Generate a list of messages. :param count: Number of messages to generate. :param order: Number of possible messages. :return: One-dimensional vector with each entry being the index of the generated message which is between 0, inclusive, and `order`, exclusive. >>> get_random_messages(5) torch.tensor([0, 2, 0, 3, 4]) """ return torch.randint(0, order, (count,)) def messages_to_onehot(messages, order): """ Convert messages represented as indexes to one-hot encoding. :param messages: List of messages to convert. :param order: Number of possible messages. :return: One-hot encoded messages. >>> messages_to_onehot(torch.tensor([0, 2, 0, 3, 4])) torch.tensor([ [1., 0., 0., 0., 0.], [0., 0., 1., 0., 0.], [1., 0., 0., 0., 0.], [0., 0., 0., 1., 0.], [0., 0., 0., 0., 1.], ]) """ return torch.nn.functional.one_hot(messages, num_classes=order).float() def plot_constellation( ax, constellation, channel, decoder, grid_step=0.05, noise_samples=1000 ): """ Plot a constellation with its decoder and channel noise. :param ax: Matplotlib axes to plot on. :param constellation: Constellation to plot. :param channel: Channel model to use for generating noise. :param decoder: Decoder function able to map the constellation points back to the original messages. :param grid_step: Grid step used for drawing the decision regions, expressed as percentage of the total plot width (or equivalently height). Lower steps makes more precise grids but takes more time to compute. :param noise_samples: Number of noisy points to sample and plot. """ ax.grid() order = len(constellation) color_map = ListedColormap(seaborn.color_palette('husl', n_colors=order)) color_norm = matplotlib.colors.BoundaryNorm(range(order + 1), order) # Extend axes symmetrically around zero so that they fit data axis_extent = max( abs(constellation.min()), abs(constellation.max()) ) * 1.05 ax.set_xlim(-axis_extent, axis_extent) ax.set_ylim(-axis_extent, axis_extent) # Hide borders but keep ticks for direction in ['left', 'bottom', 'right', 'top']: ax.axis[direction].line.set_color('#00000000') # Show zero-centered axes without ticks for direction in ['xzero', 'yzero']: axis = ax.axis[direction] axis.set_visible(True) axis.set_axisline_style('-|>') axis.major_ticklabels.set_visible(False) # Add axis names ax.annotate( 'I', (1, 0.5), xycoords='axes fraction', xytext=(25, 0), textcoords='offset points', va='center', ha='right' ) ax.annotate( 'Q', (0.5, 1), xycoords='axes fraction', xytext=(0, 25), textcoords='offset points', va='center', ha='center' ) # Plot decision regions regions_extent = 2 * axis_extent step = grid_step * regions_extent grid_range = torch.arange(-regions_extent, regions_extent, step) grid_y, grid_x = torch.meshgrid(grid_range, grid_range) grid_points = torch.stack((grid_x, grid_y), dim=-1).flatten(end_dim=1) grid_images = decoder(grid_points).argmax(dim=-1).reshape(grid_x.shape) ax.imshow( grid_images, extent=( -regions_extent, regions_extent, -regions_extent, regions_extent ), aspect='auto', origin='lower', cmap=color_map, norm=color_norm, alpha=0.2 ) # Plot constellation ax.scatter( *zip(*constellation.tolist()), zorder=10, s=60, c=range(order), edgecolor='black', cmap=color_map, norm=color_norm, ) # Plot constellation center center = constellation.sum(dim=0) / order ax.scatter( center[0], center[1], marker='X', ) # Plot channel noise if noise_samples > 0: noisy_vectors = channel(constellation.repeat(noise_samples, 1)) ax.scatter( *zip(*noisy_vectors.tolist()), marker='.', s=5, c=list(range(order)) * noise_samples, cmap=color_map, norm=color_norm, alpha=0.3, zorder=8 )