import constellation from constellation import util import torch from matplotlib import pyplot import matplotlib from mpl_toolkits.axisartist.axislines import SubplotZero # Number learned symbols order = 4 # File in which the trained model is saved input_file = 'output/constellation-order-{}.pth'.format(order) # Color map used for decision regions color_map = matplotlib.cm.Set1 # Restore model from file model = constellation.ConstellationNet( order=order, encoder_layers_sizes=(4,), decoder_layers_sizes=(4,), channel_model=constellation.GaussianChannel() ) model.load_state_dict(torch.load(input_file)) model.eval() # Extract encoding with torch.no_grad(): encoded_vectors = model.encoder( util.messages_to_onehot( torch.arange(0, order), order ) ) # Setup plot fig = pyplot.figure() ax = SubplotZero(fig, 111) fig.add_subplot(ax) # Extend axes symmetrically around zero so that they fit data extent = max( abs(encoded_vectors.min()), abs(encoded_vectors.max()) ) * 1.05 ax.set_xlim(-extent, extent) ax.set_ylim(-extent, extent) # Hide borders for direction in ['left', 'bottom', 'right', 'top']: ax.axis[direction].set_visible(False) # Show zero-centered axes for direction in ['xzero', 'yzero']: axis = ax.axis[direction] axis.set_visible(True) axis.set_axisline_style("-|>") # Configure axes ticks and labels ax.annotate( 'I', (1, 0.5), xycoords='axes fraction', xytext=(25, 0), textcoords='offset points', va='center', ha='right' ) ax.axis['xzero'].major_ticklabels.set_backgroundcolor('white') ax.axis['xzero'].major_ticklabels.set_ha('center') ax.axis['xzero'].major_ticklabels.set_va('top') ax.annotate( 'Q', (0.5, 1), xycoords='axes fraction', xytext=(0, 25), textcoords='offset points', va='center', ha='center' ) ax.axis['yzero'].major_ticklabels.set_rotation(-90) ax.axis['yzero'].major_ticklabels.set_backgroundcolor('white') ax.axis['yzero'].major_ticklabels.set_ha('left') ax.axis['yzero'].major_ticklabels.set_va('center') # Add a single tick on 0 ax.set_xticks(ax.get_xticks()[ax.get_xticks() != 0]) ax.set_yticks(ax.get_yticks()[ax.get_yticks() != 0]) ax.annotate( '0', (0, 0), xytext=(15, -10), textcoords='offset points', va='center', ha='center' ) ax.grid() # Plot decision regions color_norm = matplotlib.colors.BoundaryNorm(range(order + 1), order) step = 0.001 * extent grid_range = torch.arange(-extent, extent, step) grid_y, grid_x = torch.meshgrid(grid_range, grid_range) grid_images = model.decoder(torch.stack((grid_x, grid_y), dim=2)).argmax(dim=2) ax.imshow( grid_images, extent=(-extent, extent, -extent, extent), aspect="auto", origin="lower", cmap=color_map, norm=color_norm, alpha=0.15 ) # Plot encoded vectors ax.scatter( *zip(*encoded_vectors.tolist()), zorder=10, c=range(len(encoded_vectors)), edgecolor='black', cmap=color_map, norm=color_norm, ) for row in range(order): ax.annotate( row + 1, encoded_vectors[row], xytext=(5, 5), textcoords='offset points', backgroundcolor='white', zorder=9 ) pyplot.show()