156 lines
4.4 KiB
Python
156 lines
4.4 KiB
Python
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
|
|
)
|