Plot the two-dimensional neuron assignments. :param assignments: Vector of neuron label assignments. :param im: Used for re-drawing the assignments plot. :param figsize: Horizontal, vertical figure size in inches. :param classes: Iterable of labels for colorbar ticks correspond
(
assignments: torch.Tensor,
im: Optional[AxesImage] = None,
figsize: Tuple[int, int] = (5, 5),
classes: Optional[Sized] = None,
save: Optional[str] = None,
)
| 485 | |
| 486 | |
| 487 | def plot_assignments( |
| 488 | assignments: torch.Tensor, |
| 489 | im: Optional[AxesImage] = None, |
| 490 | figsize: Tuple[int, int] = (5, 5), |
| 491 | classes: Optional[Sized] = None, |
| 492 | save: Optional[str] = None, |
| 493 | ) -> AxesImage: |
| 494 | # language=rst |
| 495 | """ |
| 496 | Plot the two-dimensional neuron assignments. |
| 497 | |
| 498 | :param assignments: Vector of neuron label assignments. |
| 499 | :param im: Used for re-drawing the assignments plot. |
| 500 | :param figsize: Horizontal, vertical figure size in inches. |
| 501 | :param classes: Iterable of labels for colorbar ticks corresponding to data labels. |
| 502 | :param save: file name to save fig, if None = not saving fig. |
| 503 | :return: Used for re-drawing the assigments plot. |
| 504 | """ |
| 505 | locals_assignments = assignments.detach().clone().cpu().numpy() |
| 506 | |
| 507 | if save is not None: |
| 508 | plt.ioff() |
| 509 | |
| 510 | a = save.split(".") |
| 511 | if len(a) == 2: |
| 512 | save = a[0] + ".1." + a[1] |
| 513 | else: |
| 514 | a[1] = "." + str(1 + int(a[1])) + ".png" |
| 515 | save = a[0] + a[1] |
| 516 | |
| 517 | fig, ax = plt.subplots(figsize=figsize) |
| 518 | ax.set_title("Categorical assignments") |
| 519 | |
| 520 | if classes is None: |
| 521 | color = plt.get_cmap("RdBu", 11) |
| 522 | im = ax.matshow(locals_assignments, cmap=color, vmin=-1.5, vmax=9.5) |
| 523 | else: |
| 524 | color = plt.get_cmap("RdBu", len(classes) + 1) |
| 525 | im = ax.matshow( |
| 526 | locals_assignments, cmap=color, vmin=-1.5, vmax=len(classes) - 0.5 |
| 527 | ) |
| 528 | |
| 529 | div = make_axes_locatable(ax) |
| 530 | cax = div.append_axes("right", size="5%", pad=0.05) |
| 531 | |
| 532 | if classes is None: |
| 533 | cbar = plt.colorbar(im, cax=cax, ticks=list(range(-1, 11))) |
| 534 | cbar.ax.set_yticklabels(["none"] + list(range(11))) |
| 535 | else: |
| 536 | cbar = plt.colorbar(im, cax=cax, ticks=np.arange(-1, len(classes))) |
| 537 | cbar.ax.set_yticklabels(["none"] + list(classes)) |
| 538 | |
| 539 | ax.set_xticks(()) |
| 540 | ax.set_yticks(()) |
| 541 | # fig.tight_layout() |
| 542 | |
| 543 | fig.savefig(save, bbox_inches="tight") |
| 544 | plt.close() |
no test coverage detected