mlspm.visualization#

mlspm.visualization.make_input_plots(Xs: list[ndarray], outdir: str = './predictions/', start_ind: int = 0, constant_range: bool = False, cmap: str | Colormap = 'afmhot', verbose: int = 1)[source]#

Plot a batch of AFM images to files 0_input.png, 1_input.png, … etc.

Parameters:

  • Xs – Input AFM images to plot. Each list element corresponds to one AFM tip and is an array of shape (batch, x, y, z).

  • outdir – Directory where images are saved.

  • start_ind – Starting index for file naming.

  • constant_range – Whether the different slices should use the same value range or not.

  • cmap – Colormap to use for plotting.

  • verbose – Whether to print output information.

mlspm.visualization.plot_confusion_matrix(ax: Axes, conf_mat: ndarray, tick_labels: list[str] | None = None)[source]#

Plot confusion matrix on matplotlib axes.

Parameters:

  • ax – Axes object on which the confusion matrix is plotted.

  • conf_mat – Confusion matrix counts.

  • tick_labels – Labels for classes.

mlspm.visualization.plot_input(X: ndarray, constant_range: bool = False, cmap: str | Colormap = 'afmhot') → Figure[source]#

Plot a single stack of AFM images.

Parameters:

  • X – AFM image to plot.

  • constant_range – Whether the different slices should use the same value range or not.

  • cmap – Colormap to use for plotting.

Returns:

Figure on which the image was plotted.

mlspm.visualization.make_prediction_plots(preds: List[ndarray] = None, true: List[ndarray] = None, losses: ndarray = None, descriptors: List[str] = None, outdir: PathLike = './predictions/', start_ind: int = 0, verbose: bool = True)[source]#

Plot predictions/references for image descriptors.

Parameters:

  • preds – Predicted maps. Each list element corresponds to one descriptor and is an array of shape (batch_size, x_dim, y_dim).

  • true – Reference maps. Each list element corresponds to one descriptor and is an array of shape (batch_size, x_dim, y_dim).

  • losses – Losses for each prediction. Array of shape (len(preds), batch_size).

  • descriptors – Names of descriptors. The name "ES" causes the coolwarm colormap to be used.

  • outdir – Directory where images are saved.

  • start_ind – Starting index for saved images.

  • verbose – Whether to print output information.

mlspm.visualization.plot_distribution_grid(pred_dist: ndarray, ref_dist: ndarray | None = None, box_borders: ndarray = array([[2., 2., -1.5], [18., 18., 0.]]), outdir: str = './graphs/', start_ind: int = 0, verbose: int = 1)[source]#

Plot batch of position distribution grids into files 0_pred_dist.png, 1_pred_dist.png, …, and 0_pred_dist2D.png 1_pred_dist2D.png, … etc.

The full grids are divided into separate images for each z-slice in the arrays. The 2D grids are averaged over the z-dimension of the full grids.

Parameters:

  • pred_dist – Predicted position distribution grid.

  • ref_dist – Reference position distribution grid.

  • box_borders – Real-space extent of the distribution grid region in Ångströms. The array should be of the form ((x_start, y_start, z_start), (x_end, y_end, z_end)).

  • outdir – Directory where files are saved.

  • start_ind – Starting index for file naming.

  • verbose – Whether to print output information.

mlspm.visualization.plot_graphs(pred: list[MoleculeGraph] | None = None, ref: list[MoleculeGraph] | None = None, box_borders: ndarray = array([[0., 0., -1.4], [16., 16., 0.5]]), outdir: str = './graphs/', classes: list[list[int]] = None, class_colors: list[str] = 'rkbgcmy', start_ind: int = 0, verbose: int = 1)[source]#

Plot batch of graphs into files 0_graph.png, 1_graph.png, … etc.

Parameters:

  • pred – Predicted molecule graphs.

  • ref – Reference molecule graphs.

  • box_borders – Real-space extent of the plotting region in Ångströms. The array should be of the form ((x_start, y_start, z_start), (x_end, y_end, z_end)).

  • outdir – Directory where files are saved.

  • classes – Classes for categorizing atoms based on their chemical elements. Each class is a list of elements either as atomic numbers or as chemical symbols.

  • class_colors – Colors for each atom class.

  • start_ind – Starting index for file naming.

  • verbose – Whether to print output information.