gnn_tracking.analysis.latent

gnn_tracking.analysis.latent#

Plotting functions to plot the latent space

Classes#

SelectedPidsPlot

Plot the condensation space with selected PIDs highlighted.

Functions#

get_color_mapper(→ Callable[[numpy.ndarray], ...)

Get a function that maps values to colors.

Module Contents#

gnn_tracking.analysis.latent.get_color_mapper(selected_values: Sequence, colors: Sequence | None = None) → Callable[[numpy.ndarray], numpy.ndarray]#

Get a function that maps values to colors.

class gnn_tracking.analysis.latent.SelectedPidsPlot(*, condensation_space: torch.Tensor, particle_id: torch.Tensor, labels: torch.Tensor, selected_pids: Sequence[int] | None = None, ec_hit_mask: torch.Tensor, input_node_features: torch.Tensor)#

Plot the condensation space with selected PIDs highlighted. Two kinds of plots are supported: Latent space coordinates and phi/eta. For each of these, separate methods plot hits of the selected PIDs, all other hits, and collateral hits (hits in the same cluster as the selected PIDs).

Parameters:

  • condensation_space

  • particle_id

  • labels

  • selected_pids

  • ec_hit_mask – If we do orphan node prediction, we need to know which hits make it to the condensation space

  • input_node_features

_ec_hit_mask#
_x#
_pids#
_labels#
_selected_pids#
_color_mapper#
_selected_pid_mask#
_phi#
_eta#
get_collateral_mask(pid: int) → torch.Tensor#

Mask for hits that are in the same cluster(s) as the hits belonging to this particle ID.

static plot_circles(ax: matplotlib.pyplot.Axes, xs: torch.Tensor, ys: torch.Tensor, colors, eps=1) → None#
get_colors(pids: torch.Tensor | Sequence) → Sequence#
plot_selected_pid_latent(ax: matplotlib.pyplot.Axes, plot_circles=False) → None#
plot_collateral_latent(ax: matplotlib.pyplot.Axes) → None#
plot_other_hit_latent(ax: matplotlib.pyplot.Axes) → None#
plot_selected_pid_ep(ax: matplotlib.pyplot.Axes) → None#
plot_other_hit_ep(ax: matplotlib.pyplot.Axes) → None#
plot_collateral_ep(ax: matplotlib.pyplot.Axes) → None#
Contents