Function plot_edge_affinity_distributions

deeplabcut/utils/plotting.py:437–518

Display the distribution of affinity costs of within- and between-animal edges. Parameters ---------- eval_pickle_file : string Path to a *_full.pickle from the evaluation-results folder. include_bodyparts : list of strings, optional A list of body part names whose

(
    eval_pickle_file,
    include_bodyparts="all",
    output_name="",
    figsize=(10, 7),
)

def plot_edge_affinity_distributions(
    eval_pickle_file,
    include_bodyparts="all",
    output_name="",
    figsize=(10, 7),
):
    """Display the distribution of affinity costs of within- and between-animal edges.

    Parameters
    ----------
    eval_pickle_file : string
        Path to a *_full.pickle from the evaluation-results folder.

    include_bodyparts : list of strings, optional
        A list of body part names whose edges are to be shown.
        By default, all body parts and their corresponding edges are analyzed.
        We recommend only passing a subset of body parts for projects with large graphs.

    output_name: string, optional
        Path where the plot is saved. By default, it is stored as costdist.png.

    figsize: tuple
        Figure size in inches.
    """

    with open(eval_pickle_file, "rb") as file:
        data = pickle.load(file)
    meta_pickle_file = eval_pickle_file.replace("_full.", "_meta.")
    with open(meta_pickle_file, "rb") as file:
        metadata = pickle.load(file)
    (w_train, _), (b_train, _) = crossvalutils._calc_within_between_pafs(
        data,
        metadata,
        train_set_only=True,
    )
    data.pop("metadata", None)
    nonempty = set(i for i, vals in w_train.items() if vals)
    meta = metadata["data"]["DLC-model-config file"]
    bpts = list(map(str.lower, meta["all_joints_names"]))
    inds_multi = set(b for edge in meta["partaffinityfield_graph"] for b in edge)
    if include_bodyparts == "all":
        include_bodyparts = inds_multi
    else:
        include_bodyparts = set(bpts.index(bpt) for bpt in include_bodyparts)
    edges_to_keep = set()
    graph = meta["partaffinityfield_graph"]
    for n, edge in enumerate(graph):
        if not any(i in include_bodyparts for i in edge):
            continue
        edges_to_keep.add(n)
    edge_inds = edges_to_keep.intersection(nonempty)
    nrows = int(np.ceil(np.sqrt(len(edge_inds))))
    ncols = int(np.ceil(len(edge_inds) / nrows))
    fig, axes_ = plt.subplots(
        nrows,
        ncols,
        figsize=figsize,
        tight_layout=True,