146 lines
3.5 KiB
Python
146 lines
3.5 KiB
Python
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
|
|
axis_extent = max(
|
|
abs(encoded_vectors.min()),
|
|
abs(encoded_vectors.max())
|
|
) * 1.05
|
|
ax.set_xlim(-axis_extent, axis_extent)
|
|
ax.set_ylim(-axis_extent, axis_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'].set_zorder(9)
|
|
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'].set_zorder(9)
|
|
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)
|
|
|
|
regions_extent = 2 * axis_extent
|
|
step = 0.001 * regions_extent
|
|
grid_range = torch.arange(-regions_extent, regions_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=(-regions_extent, regions_extent, -regions_extent, regions_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
|
|
)
|
|
|
|
# Plot noise
|
|
noisy_count = 1000
|
|
noisy_vectors = model.channel(encoded_vectors.repeat(noisy_count, 1))
|
|
|
|
ax.scatter(
|
|
*zip(*noisy_vectors.tolist()),
|
|
s=1,
|
|
c=list(range(len(encoded_vectors))) * noisy_count,
|
|
cmap=color_map,
|
|
norm=color_norm,
|
|
zorder=8
|
|
)
|
|
|
|
pyplot.show()
|