--- a
+++ b/src/nichecompass/utils/analysis.py
@@ -0,0 +1,1147 @@
+"""
+This module contains utilities to analyze niches inferred by the NicheCompass
+model.
+"""
+
+from typing import Optional, Tuple
+
+#import holoviews as hv
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import scanpy as sc
+import scipy.sparse as sp
+import seaborn as sns
+from anndata import AnnData
+from matplotlib import cm, colors
+from matplotlib.lines import Line2D
+import networkx as nx
+
+from ..models import NicheCompass
+
+
+def aggregate_obsp_matrix_per_cell_type(
+        adata: AnnData,
+        obsp_key: str,
+        cell_type_key: str="cell_type",
+        group_key: Optional[str]=None,
+        agg_rows: bool=False):
+    """
+    Generic function to aggregate adjacency matrices stored in
+    ´adata.obsp[obsp_key]´ on cell type level. It can be used to aggregate the
+    node label aggregator aggregation weights alpha or the reconstructed adjacency
+    matrix of a trained NicheCompass model by neighbor cell type for downstream
+    analysis.
+
+    Parameters
+    ----------
+    adata:
+        AnnData object which contains outputs of NicheCompass model training.
+    obsp_key:
+        Key in ´adata.obsp´ where the matrix to be aggregated is stored.
+    cell_type_key:
+        Key in ´adata.obs´ where the cell type labels are stored.
+    group_key:
+        Key in ´adata.obs´ where additional grouping labels are stored.    
+    agg_rows:
+        If ´True´, also aggregate over the observations on cell type level.
+
+    Returns
+    ----------
+    cell_type_agg_df:
+        Pandas DataFrame with the aggregated obsp values (dim: n_obs x
+        n_cell_types if ´agg_rows == False´, else n_cell_types x n_cell_types).
+    """
+    n_obs = len(adata)
+    n_cell_types = adata.obs[cell_type_key].nunique()
+    sorted_cell_types = sorted(adata.obs[cell_type_key].unique().tolist())
+
+    cell_type_label_encoder = {k: v for k, v in zip(
+        sorted_cell_types,
+        range(n_cell_types))}
+
+    # Retrieve non zero indices and non zero values, and create row-wise
+    # observation cell type index
+    nz_obsp_idx = adata.obsp[obsp_key].nonzero()
+    neighbor_cell_type_index = adata.obs[cell_type_key][nz_obsp_idx[1]].map(
+        cell_type_label_encoder).values
+    adata.obsp[obsp_key].eliminate_zeros() # In some sparse reps 0s can appear
+    nz_obsp = adata.obsp[obsp_key].data
+
+    # Use non zero indices, non zero values and row-wise observation cell type
+    # index to construct new df with cell types as columns and row-wise
+    # aggregated values per cell type index as values
+    cell_type_agg = np.zeros((n_obs, n_cell_types))
+    np.add.at(cell_type_agg,
+              (nz_obsp_idx[0], neighbor_cell_type_index),
+              nz_obsp)
+    cell_type_agg_df = pd.DataFrame(
+        cell_type_agg,
+        columns=sorted_cell_types)
+    
+    # Add cell type labels of observations
+    cell_type_agg_df[cell_type_key] = adata.obs[cell_type_key].values
+
+    # If specified, add group label
+    if group_key is not None:
+        cell_type_agg_df[group_key] = adata.obs[group_key].values
+
+    if agg_rows:
+        # In addition, aggregate values across rows to get a
+        # (n_cell_types x n_cell_types) df
+        if group_key is not None:
+            cell_type_agg_df = cell_type_agg_df.groupby(
+                [group_key, cell_type_key]).sum()
+        else:
+            cell_type_agg_df = cell_type_agg_df.groupby(cell_type_key).sum()
+
+        # Sort index to have same order as columns
+        cell_type_agg_df = cell_type_agg_df.loc[
+            sorted(cell_type_agg_df.index.tolist()), :]
+        
+    return cell_type_agg_df
+
+
+def create_cell_type_chord_plot_from_df(
+        adata: AnnData,
+        df: pd.DataFrame,
+        link_threshold: float=0.01,
+        cell_type_key: str="cell_type",
+        group_key: Optional[str]=None,
+        groups: str="all",
+        plot_label: str="Niche",
+        save_fig: bool=False,
+        file_path: Optional[str]=None):
+    """
+    Create a cell type chord diagram per group based on an input DataFrame.
+
+    Parameters
+    ----------
+    adata:
+        AnnData object which contains outputs of NicheCompass model training.
+    df:
+        A Pandas DataFrame that contains the connection values for the chord
+        plot (dim: (n_groups x n_cell_types) x n_cell_types).
+    link_threshold:
+        Ratio of link strength that a cell type pair needs to exceed compared to
+        the cell type pair with the maximum link strength to be considered a
+        link for the chord plot.
+    cell_type_key:
+        Key in ´adata.obs´ where the cell type labels are stored.
+    group_key:
+        Key in ´adata.obs´ where additional group labels are stored.
+    groups:
+        List of groups that will be plotted. If ´all´, plot all groups.
+    plot_label:
+        Shared label for the plots.
+    save_fig:
+        If ´True´, save the figure.
+    file_path:
+        Path where to save the figure.
+    """
+    hv.extension("bokeh")
+    hv.output(size=200)
+
+    sorted_cell_types = sorted(adata.obs[cell_type_key].unique().tolist())
+
+    # Get group labels
+    if (group_key is not None) & (groups == "all"):
+        group_labels = df.index.get_level_values(
+            df.index.names.index(group_key)).unique().tolist()
+    elif (group_key is not None) & (groups != "all"):
+        group_labels = groups
+    else:
+        group_labels = [""]
+
+    chord_list = []
+    for group_label in group_labels:
+        if group_label == "":
+            group_df = df
+        else:
+            group_df = df[df.index.get_level_values(
+                df.index.names.index(group_key)) == group_label]
+        
+        # Get max value (over rows and columns) of the group for thresholding
+        group_max = group_df.max().max()
+
+        # Create group chord links
+        links_list = []
+        for i in range(len(sorted_cell_types)):
+            for j in range(len(sorted_cell_types)):
+                if group_df.iloc[i, j] > group_max * link_threshold:
+                    link_dict = {}
+                    link_dict["source"] = j
+                    link_dict["target"] = i
+                    link_dict["value"] = group_df.iloc[i, j]
+                    links_list.append(link_dict)
+        links = pd.DataFrame(links_list)
+
+        # Create group chord nodes (only where links exist)
+        nodes_list = []
+        nodes_idx = []
+        for i, cell_type in enumerate(sorted_cell_types):
+            if i in (links["source"].values) or i in (links["target"].values):
+                nodes_idx.append(i)
+                nodes_dict = {}
+                nodes_dict["name"] = cell_type
+                nodes_dict["group"] = 1
+                nodes_list.append(nodes_dict)
+        nodes = hv.Dataset(pd.DataFrame(nodes_list, index=nodes_idx), "index")
+
+        # Create group chord plot
+        chord = hv.Chord((links, nodes)).select(value=(5, None))
+        chord.opts(hv.opts.Chord(cmap="Category20",
+                                 edge_cmap="Category20",
+                                 edge_color=hv.dim("source").str(),
+                                 labels="name",
+                                 node_color=hv.dim("index").str(),
+                                 title=f"{plot_label} {group_label}"))
+        chord_list.append(chord)
+    
+    # Display chord plots
+    layout = hv.Layout(chord_list).cols(2)
+    hv.output(layout)
+
+    # Save chord plots
+    if save_fig:
+        hv.save(layout,
+                file_path,
+                fmt="png")
+
+        
+def generate_enriched_gp_info_plots(plot_label: str,
+                                    model: NicheCompass,
+                                    sample_key: str,
+                                    differential_gp_test_results_key: str,
+                                    cat_key: str,
+                                    cat_palette: dict,
+                                    n_top_enriched_gp_start_idx: int=0,
+                                    n_top_enriched_gp_end_idx: int=10,
+                                    feature_spaces: list=["latent"],
+                                    n_top_genes_per_gp: int=3,
+                                    n_top_peaks_per_gp: int=0,
+                                    scale_omics_ft: bool=False,
+                                    save_figs: bool=False,
+                                    figure_folder_path: str="",
+                                    file_format: str="png",
+                                    spot_size: float=30.):
+    """
+    Generate info plots of enriched gene programs. These show the enriched
+    category, the gp activities, as well as the counts (or log normalized
+    counts) of the top genes and/or peaks in a specified feature space.
+    
+    Parameters
+    ----------
+    plot_label:
+        Main label of the plots.
+    model:
+        A trained NicheCompass model.
+    sample_key:
+        Key in ´adata.obs´ where the samples are stored.
+    differential_gp_test_results_key:
+        Key in ´adata.uns´ where the results of the differential gene program
+        testing are stored.
+    cat_key:
+        Key in ´adata.obs´ where the categories that are used as colors for the
+        enriched category plot are stored.
+    cat_palette:
+        Dictionary of colors that are used to highlight the categories, where
+        the category is the key of the dictionary and the color is the value.
+    n_top_enriched_gp_start_idx:
+        Number of top enriched gene program from which to start the creation
+        of plots.
+    n_top_enriched_gp_end_idx:
+        Number of top enriched gene program at which to stop the creation
+        of plots.
+    feature_spaces:
+        List of feature spaces used for the info plots. Can be ´latent´ to use
+        the latent embeddings for the plots, or it can be any of the samples
+        stored in ´adata.obs[sample_key]´ to use the respective physical
+        feature space for the plots.
+    n_top_genes_per_gp:
+        Number of top genes per gp to be considered in the info plots.
+    n_top_peaks_per_gp:
+        Number of top peaks per gp to be considered in the info plots. If ´>0´,
+        requires the model to be trained inlcuding ATAC modality.
+    scale_omics_ft:
+        If ´True´, scale genes and peaks before plotting.
+    save_figs:
+        If ´True´, save the figures.
+    figure_folder_path:
+        Folder path where the figures will be saved.
+    file_format:
+        Format with which the figures will be saved.
+    spot_size:
+        Spot size used for the spatial plots.
+    """
+    model._check_if_trained(warn=True)
+
+    adata = model.adata.copy()
+    if n_top_peaks_per_gp > 0:
+        if "atac" not in model.modalities_:
+            raise ValueError("The model needs to be trained with ATAC data if"
+                             "'n_top_peaks_per_gp' > 0.")
+        adata_atac = model.adata_atac.copy()
+    
+    # TODO
+    if scale_omics_ft:
+        sc.pp.scale(adata)
+        if n_top_peaks_per_gp > 0:
+            sc.pp.scale(adata_atac)
+        adata.uns["omics_ft_pos_cmap"] = "RdBu"
+        adata.uns["omics_ft_neg_cmap"] = "RdBu_r"
+    else:
+        if n_top_peaks_per_gp > 0:
+            adata_atac.X = adata_atac.X.toarray()
+        adata.uns["omics_ft_pos_cmap"] = "Blues"
+        adata.uns["omics_ft_neg_cmap"] = "Reds"
+        
+    cats = list(adata.uns[differential_gp_test_results_key]["category"][
+        n_top_enriched_gp_start_idx:n_top_enriched_gp_end_idx])
+    gps = list(adata.uns[differential_gp_test_results_key]["gene_program"][
+        n_top_enriched_gp_start_idx:n_top_enriched_gp_end_idx])
+    log_bayes_factors = list(adata.uns[differential_gp_test_results_key]["log_bayes_factor"][
+        n_top_enriched_gp_start_idx:n_top_enriched_gp_end_idx])
+    
+    for gp in gps:
+        # Get source and target genes, gene importances and gene signs and store
+        # in temporary adata
+        gp_gene_importances_df = model.compute_gp_gene_importances(
+            selected_gp=gp)
+        
+        gp_source_genes_gene_importances_df = gp_gene_importances_df[
+            gp_gene_importances_df["gene_entity"] == "source"]
+        gp_target_genes_gene_importances_df = gp_gene_importances_df[
+            gp_gene_importances_df["gene_entity"] == "target"]
+        adata.uns["n_top_source_genes"] = n_top_genes_per_gp
+        adata.uns[f"{gp}_source_genes_top_genes"] = (
+            gp_source_genes_gene_importances_df["gene"][
+                :n_top_genes_per_gp].values)
+        adata.uns[f"{gp}_source_genes_top_gene_importances"] = (
+            gp_source_genes_gene_importances_df["gene_importance"][
+                :n_top_genes_per_gp].values)
+        adata.uns[f"{gp}_source_genes_top_gene_signs"] = (
+            np.where(gp_source_genes_gene_importances_df[
+                "gene_weight"] > 0, "+", "-"))
+        adata.uns["n_top_target_genes"] = n_top_genes_per_gp
+        adata.uns[f"{gp}_target_genes_top_genes"] = (
+            gp_target_genes_gene_importances_df["gene"][
+                :n_top_genes_per_gp].values)
+        adata.uns[f"{gp}_target_genes_top_gene_importances"] = (
+            gp_target_genes_gene_importances_df["gene_importance"][
+                :n_top_genes_per_gp].values)
+        adata.uns[f"{gp}_target_genes_top_gene_signs"] = (
+            np.where(gp_target_genes_gene_importances_df[
+                "gene_weight"] > 0, "+", "-"))
+
+        if n_top_peaks_per_gp > 0:
+            # Get source and target peaks, peak importances and peak signs and
+            # store in temporary adata
+            gp_peak_importances_df = model.compute_gp_peak_importances(
+                selected_gp=gp)
+            gp_source_peaks_peak_importances_df = gp_peak_importances_df[
+                gp_peak_importances_df["peak_entity"] == "source"]
+            gp_target_peaks_peak_importances_df = gp_peak_importances_df[
+                gp_peak_importances_df["peak_entity"] == "target"]
+            adata.uns["n_top_source_peaks"] = n_top_peaks_per_gp
+            adata.uns[f"{gp}_source_peaks_top_peaks"] = (
+                gp_source_peaks_peak_importances_df["peak"][
+                    :n_top_peaks_per_gp].values)
+            adata.uns[f"{gp}_source_peaks_top_peak_importances"] = (
+                gp_source_peaks_peak_importances_df["peak_importance"][
+                    :n_top_peaks_per_gp].values)
+            adata.uns[f"{gp}_source_peaks_top_peak_signs"] = (
+                np.where(gp_source_peaks_peak_importances_df[
+                    "peak_weight"] > 0, "+", "-"))
+            adata.uns["n_top_target_peaks"] = n_top_peaks_per_gp
+            adata.uns[f"{gp}_target_peaks_top_peaks"] = (
+                gp_target_peaks_peak_importances_df["peak"][
+                    :n_top_peaks_per_gp].values)
+            adata.uns[f"{gp}_target_peaks_top_peak_importances"] = (
+                gp_target_peaks_peak_importances_df["peak_importance"][
+                    :n_top_peaks_per_gp].values)
+            adata.uns[f"{gp}_target_peaks_top_peak_signs"] = (
+                np.where(gp_target_peaks_peak_importances_df[
+                    "peak_weight"] > 0, "+", "-"))
+            
+            # Add peak counts to temporary adata for plotting
+            adata.obs[[peak for peak in 
+                       adata.uns[f"{gp}_target_peaks_top_peaks"]]] = (
+                adata_atac.X[
+                    :, [adata_atac.var_names.tolist().index(peak)
+                        for peak in adata.uns[f"{gp}_target_peaks_top_peaks"]]])
+            adata.obs[[peak for peak in
+                       adata.uns[f"{gp}_source_peaks_top_peaks"]]] = (
+                adata_atac.X[
+                    :, [adata_atac.var_names.tolist().index(peak)
+                        for peak in adata.uns[f"{gp}_source_peaks_top_peaks"]]])
+        else:
+            adata.uns["n_top_source_peaks"] = 0
+            adata.uns["n_top_target_peaks"] = 0
+
+    for feature_space in feature_spaces:
+        plot_enriched_gp_info_plots_(
+            adata=adata,
+            sample_key=sample_key,
+            gps=gps,
+            log_bayes_factors=log_bayes_factors,
+            cat_key=cat_key,
+            cat_palette=cat_palette,
+            cats=cats,
+            feature_space=feature_space,
+            spot_size=spot_size,
+            suptitle=f"{plot_label.replace('_', ' ').title()} "
+                     f"Top {n_top_enriched_gp_start_idx} to "
+                     f"{n_top_enriched_gp_end_idx} Enriched GPs: "
+                     f"GP Scores and Omics Feature Counts in "
+                     f"{feature_space} Feature Space",
+            save_fig=save_figs,
+            figure_folder_path=figure_folder_path,
+            fig_name=f"{plot_label}_top_{n_top_enriched_gp_start_idx}"
+                     f"-{n_top_enriched_gp_end_idx}_enriched_gps_gp_scores_"
+                     f"omics_feature_counts_in_{feature_space}_"
+                     f"feature_space.{file_format}")
+            
+            
+def plot_enriched_gp_info_plots_(adata: AnnData,
+                                 sample_key: str,
+                                 gps: list,
+                                 log_bayes_factors: list,
+                                 cat_key: str,
+                                 cat_palette: dict,
+                                 cats: list,
+                                 feature_space: str,
+                                 spot_size: float,
+                                 suptitle: str,
+                                 save_fig: bool,
+                                 figure_folder_path: str,
+                                 fig_name: str):
+    """
+    This is a helper function to plot gene program info plots in a specified
+    feature space.
+    
+    Parameters
+    ----------
+    adata:
+        An AnnData object with stored information about the gene programs to be
+        plotted.
+    sample_key:
+        Key in ´adata.obs´ where the samples are stored.
+    gps:
+        List of gene programs for which info plots will be created.
+    log_bayes_factors:
+        List of log bayes factors corresponding to gene programs
+    cat_key:
+        Key in ´adata.obs´ where the categories that are used as colors for the
+        enriched category plot are stored.
+    cat_palette:
+        Dictionary of colors that are used to highlight the categories, where
+        the category is the key of the dictionary and the color is the value.
+    cats:
+        List of categories for which the corresponding gene programs in ´gps´
+        are enriched.
+    feature_space:
+        Feature space used for the plots. Can be ´latent´ to use the latent
+        embeddings for the plots, or it can be any of the samples stored in
+        ´adata.obs[sample_key]´ to use the respective physical feature space for
+        the plots.
+    spot_size:
+        Spot size used for the spatial plots.
+    subtitle:
+        Overall figure title.
+    save_fig:
+        If ´True´, save the figure.
+    figure_folder_path:
+        Path of the folder where the figure will be saved.
+    fig_name:
+        Name of the figure under which it will be saved.
+    """
+    # Define figure configurations
+    ncols = (2 +
+             adata.uns["n_top_source_genes"] +
+             adata.uns["n_top_target_genes"] +
+             adata.uns["n_top_source_peaks"] +
+             adata.uns["n_top_target_peaks"])
+    fig_width = (12 + (6 * (
+        adata.uns["n_top_source_genes"] +
+        adata.uns["n_top_target_genes"] +
+        adata.uns["n_top_source_peaks"] +
+        adata.uns["n_top_target_peaks"])))
+    wspace = 0.3
+    fig, axs = plt.subplots(nrows=len(gps),
+                            ncols=ncols,
+                            figsize=(fig_width, 6*len(gps)))
+    if axs.ndim == 1:
+        axs = axs.reshape(1, -1)
+    title = fig.suptitle(t=suptitle,
+                         x=0.55,
+                         y=(1.1 if len(gps) == 1 else 0.97),
+                         fontsize=20)
+    
+    # Plot enriched gp category and gene program latent scores
+    for i, gp in enumerate(gps):
+        if feature_space == "latent":
+            sc.pl.umap(
+                adata,
+                color=cat_key,
+                palette=cat_palette,
+                groups=cats[i],
+                ax=axs[i, 0],
+                title="Enriched GP Category",
+                legend_loc="on data",
+                na_in_legend=False,
+                show=False)
+            sc.pl.umap(
+                adata,
+                color=gps[i],
+                color_map="RdBu",
+                ax=axs[i, 1],
+                title=f"{gp[:gp.index('_')]}\n"
+                      f"{gp[gp.index('_') + 1: gp.rindex('_')].replace('_', ' ')}"
+                      f"\n{gp[gps[i].rindex('_') + 1:]} score (LBF: {round(log_bayes_factors[i])})",
+                colorbar_loc="bottom",
+                show=False)
+        else:
+            sc.pl.spatial(
+                adata=adata[adata.obs[sample_key] == feature_space],
+                color=cat_key,
+                palette=cat_palette,
+                groups=cats[i],
+                ax=axs[i, 0],
+                spot_size=spot_size,
+                title="Enriched GP Category",
+                legend_loc="on data",
+                na_in_legend=False,
+                show=False)
+            sc.pl.spatial(
+                adata=adata[adata.obs[sample_key] == feature_space],
+                color=gps[i],
+                color_map="RdBu",
+                spot_size=spot_size,
+                title=f"{gps[i].split('_', 1)[0]}\n{gps[i].split('_', 1)[1]} "
+                      f"(LBF: {round(log_bayes_factors[i], 2)})",
+                legend_loc=None,
+                ax=axs[i, 1],
+                colorbar_loc="bottom",
+                show=False) 
+        axs[i, 0].xaxis.label.set_visible(False)
+        axs[i, 0].yaxis.label.set_visible(False)
+        axs[i, 1].xaxis.label.set_visible(False)
+        axs[i, 1].yaxis.label.set_visible(False)
+
+        # Plot omics feature counts (or log normalized counts)
+        modality_entities = []
+        if len(adata.uns[f"{gp}_source_genes_top_genes"]) > 0:
+            modality_entities.append("source_genes")
+        if len(adata.uns[f"{gp}_target_genes_top_genes"]) > 0:
+            modality_entities.append("target_genes")
+        if f"{gp}_source_peaks_top_peaks" in adata.uns.keys():
+            gp_n_source_peaks_top_peaks = (
+                len(adata.uns[f"{gp}_source_peaks_top_peaks"]))
+            if len(adata.uns[f"{gp}_source_peaks_top_peaks"]) > 0:
+                modality_entities.append("source_peaks")
+        else:
+            gp_n_source_peaks_top_peaks = 0
+        if f"{gp}_target_peaks_top_peaks" in adata.uns.keys():
+            gp_n_target_peaks_top_peaks = (
+                len(adata.uns[f"{gp}_target_peaks_top_peaks"]))
+            if len(adata.uns[f"{gp}_target_peaks_top_peaks"]) > 0:
+                modality_entities.append("target_peaks")
+        else:
+            gp_n_target_peaks_top_peaks = 0
+        for modality_entity in modality_entities:
+            # Define k for index iteration
+            if modality_entity == "source_genes":
+                k = 0
+            elif modality_entity == "target_genes":
+                k = len(adata.uns[f"{gp}_source_genes_top_genes"])
+            elif modality_entity == "source_peaks":
+                k = (len(adata.uns[f"{gp}_source_genes_top_genes"]) +
+                     len(adata.uns[f"{gp}_target_genes_top_genes"]))
+            elif modality_entity == "target_peaks":
+                k = (len(adata.uns[f"{gp}_source_genes_top_genes"]) +
+                     len(adata.uns[f"{gp}_target_genes_top_genes"]) +
+                     len(adata.uns[f"{gp}_source_peaks_top_peaks"]))
+            for j in range(len(adata.uns[f"{gp}_{modality_entity}_top_"
+                                         f"{modality_entity.split('_')[1]}"])):
+                if feature_space == "latent":
+                    sc.pl.umap(
+                        adata,
+                        color=adata.uns[f"{gp}_{modality_entity}_top_"
+                                        f"{modality_entity.split('_')[1]}"][j],
+                        color_map=(adata.uns["omics_ft_pos_cmap"] if
+                                   adata.uns[f"{gp}_{modality_entity}_top_"
+                                             f"{modality_entity.split('_')[1][:-1]}"
+                                             "_signs"][j] == "+" else adata.uns["omics_ft_neg_cmap"]),
+                        ax=axs[i, 2+k+j],
+                        legend_loc="on data",
+                        na_in_legend=False,
+                        title=f"""{adata.uns[f"{gp}_{modality_entity}_top_"
+                                             f"{modality_entity.split('_')[1]}"
+                                             ][j]}: """
+                              f"""{adata.uns[f"{gp}_{modality_entity}_top_"
+                                             f"{modality_entity.split('_')[1][:-1]}"
+                                             "_importances"][j]:.2f} """
+                              f"({modality_entity[:-1]}; "
+                              f"""{adata.uns[f"{gp}_{modality_entity}_top_"
+                                             f"{modality_entity.split('_')[1][:-1]}"
+                                             "_signs"][j]})""",
+                        colorbar_loc="bottom",
+                        show=False)
+                else:
+                    sc.pl.spatial(
+                        adata=adata[adata.obs[sample_key] == feature_space],
+                        color=adata.uns[f"{gp}_{modality_entity}_top_"
+                                        f"{modality_entity.split('_')[1]}"][j],
+                        color_map=(adata.uns["omics_ft_pos_cmap"] if
+                                   adata.uns[f"{gp}_{modality_entity}_top_"
+                                             f"{modality_entity.split('_')[1][:-1]}"
+                                             "_signs"][j] == "+" else adata.uns["omics_ft_neg_cmap"]),
+                        legend_loc="on data",
+                        na_in_legend=False,
+                        ax=axs[i, 2+k+j],
+                        spot_size=spot_size,
+                        title=f"""{adata.uns[f"{gp}_{modality_entity}_top_"
+                                             f"{modality_entity.split('_')[1]}"
+                                             ][j]} \n"""
+                              f"""({adata.uns[f"{gp}_{modality_entity}_top_"
+                                             f"{modality_entity.split('_')[1][:-1]}"
+                                             "_importances"][j]:.2f}; """
+                              f"{modality_entity[:-1]}; "
+                              f"""{adata.uns[f"{gp}_{modality_entity}_top_"
+                                             f"{modality_entity.split('_')[1][:-1]}"
+                                             "_signs"][j]})""",
+                        colorbar_loc="bottom",
+                        show=False)
+                axs[i, 2+k+j].xaxis.label.set_visible(False)
+                axs[i, 2+k+j].yaxis.label.set_visible(False)
+            # Remove unnecessary axes
+            for l in range(2 +
+                           len(adata.uns[f"{gp}_source_genes_top_genes"]) +
+                           len(adata.uns[f"{gp}_target_genes_top_genes"]) +
+                           gp_n_source_peaks_top_peaks +
+                           gp_n_target_peaks_top_peaks, ncols):
+                axs[i, l].set_visible(False)
+
+    # Save and display plot
+    plt.subplots_adjust(wspace=wspace, hspace=0.275)
+    if save_fig:
+        fig.savefig(f"{figure_folder_path}/{fig_name}",
+                    bbox_extra_artists=(title,),
+                    bbox_inches="tight")
+    plt.show()
+
+default_color_dict = {
+    "0": "#66C5CC",
+    "1": "#F6CF71",
+    "2": "#F89C74",
+    "3": "#DCB0F2",
+    "4": "#87C55F",
+    "5": "#9EB9F3",
+    "6": "#FE88B1",
+    "7": "#C9DB74",
+    "8": "#8BE0A4",
+    "9": "#B497E7",
+    "10": "#D3B484",
+    "11": "#B3B3B3",
+    "12": "#276A8C", # Royal Blue
+    "13": "#DAB6C4", # Pink
+    "14": "#C38D9E", # Mauve-Pink
+    "15": "#9D88A2", # Mauve
+    "16": "#FF4D4D", # Light Red
+    "17": "#9B4DCA", # Lavender-Purple
+    "18": "#FF9CDA", # Bright Pink
+    "19": "#FF69B4", # Hot Pink
+    "20": "#FF00FF", # Magenta
+    "21": "#DA70D6", # Orchid
+    "22": "#BA55D3", # Medium Orchid
+    "23": "#8A2BE2", # Blue Violet
+    "24": "#9370DB", # Medium Purple
+    "25": "#7B68EE", # Medium Slate Blue
+    "26": "#4169E1", # Royal Blue
+    "27": "#FF8C8C", # Salmon Pink
+    "28": "#FFAA80", # Light Coral
+    "29": "#48D1CC", # Medium Turquoise
+    "30": "#40E0D0", # Turquoise
+    "31": "#00FF00", # Lime
+    "32": "#7FFF00", # Chartreuse
+    "33": "#ADFF2F", # Green Yellow
+    "34": "#32CD32", # Lime Green
+    "35": "#228B22", # Forest Green
+    "36": "#FFD8B8", # Peach
+    "37": "#008080", # Teal
+    "38": "#20B2AA", # Light Sea Green
+    "39": "#00FFFF", # Cyan
+    "40": "#00BFFF", # Deep Sky Blue
+    "41": "#4169E1", # Royal Blue
+    "42": "#0000CD", # Medium Blue
+    "43": "#00008B", # Dark Blue
+    "44": "#8B008B", # Dark Magenta
+    "45": "#FF1493", # Deep Pink
+    "46": "#FF4500", # Orange Red
+    "47": "#006400", # Dark Green
+    "48": "#FF6347", # Tomato
+    "49": "#FF7F50", # Coral
+    "50": "#CD5C5C", # Indian Red
+    "51": "#B22222", # Fire Brick
+    "52": "#FFB83F",  # Light Orange
+    "53": "#8B0000", # Dark Red
+    "54": "#D2691E", # Chocolate
+    "55": "#A0522D", # Sienna
+    "56": "#800000", # Maroon
+    "57": "#808080", # Gray
+    "58": "#A9A9A9", # Dark Gray
+    "59": "#C0C0C0", # Silver
+    "60": "#9DD84A",
+    "61": "#F5F5F5", # White Smoke
+    "62": "#F17171", # Light Red
+    "63": "#000000", # Black
+    "64": "#FF8C42", # Tangerine
+    "65": "#F9A11F", # Bright Orange-Yellow
+    "66": "#FACC15", # Golden Yellow
+    "67": "#E2E062", # Pale Lime
+    "68": "#BADE92", # Soft Lime
+    "69": "#70C1B3", # Greenish-Blue
+    "70": "#41B3A3", # Turquoise
+    "71": "#5EAAA8", # Gray-Green
+    "72": "#72B01D", # Chartreuse
+    "73": "#9CD08F", # Light Green
+    "74": "#8EBA43", # Olive Green
+    "75": "#FAC8C3", # Light Pink
+    "76": "#E27D60", # Dark Salmon
+    "77": "#C38D9E", # Mauve-Pink
+    "78": "#937D64", # Light Brown
+    "79": "#B1C1CC", # Light Blue-Gray
+    "80": "#88A0A8", # Gray-Blue-Green
+    "81": "#4E598C", # Dark Blue-Purple
+    "82": "#4B4E6D", # Dark Gray-Blue
+    "83": "#8E9AAF", # Light Blue-Grey
+    "84": "#C0D6DF", # Pale Blue-Grey
+    "85": "#97C1A9", # Blue-Green
+    "86": "#4C6E5D", # Dark Green
+    "87": "#95B9C7", # Pale Blue-Green
+    "88": "#C1D5E0", # Pale Gray-Blue
+    "89": "#ECDB54", # Bright Yellow
+    "90": "#E89B3B", # Bright Orange
+    "91": "#CE5A57", # Deep Red
+    "92": "#C3525A", # Dark Red
+    "93": "#B85D8E", # Berry
+    "94": "#7D5295", # Deep Purple
+    "-1" : "#E1D9D1",
+    "None" : "#E1D9D1"
+}
+
+def create_new_color_dict(
+        adata,
+        cat_key,
+        color_palette="default",
+        overwrite_color_dict={"-1" : "#E1D9D1"},
+        skip_default_colors=0):
+    """
+    Create a dictionary of color hexcodes for a specified category.
+
+    Parameters
+    ----------
+    adata:
+        AnnData object.
+    cat_key:
+        Key in ´adata.obs´ where the categories are stored for which color
+        hexcodes will be created.
+    color_palette:
+        Type of color palette.
+    overwrite_color_dict:
+        Dictionary with overwrite values that will take precedence over the
+        automatically created dictionary.
+    skip_default_colors:
+        Number of colors to skip from the default color dict.
+
+    Returns
+    ----------
+    new_color_dict:
+        The color dictionary with a hexcode for each category.
+    """
+    new_categories = adata.obs[cat_key].unique().tolist()
+    if color_palette == "cell_type_30":
+        # https://github.com/scverse/scanpy/blob/master/scanpy/plotting/palettes.py#L40
+        new_color_dict = {key: value for key, value in zip(
+            new_categories,
+            ["#023fa5",
+             "#7d87b9",
+             "#bec1d4",
+             "#d6bcc0",
+             "#bb7784",
+             "#8e063b",
+             "#4a6fe3",
+             "#8595e1",
+             "#b5bbe3",
+             "#e6afb9",
+             "#e07b91",
+             "#d33f6a",
+             "#11c638",
+             "#8dd593",
+             "#c6dec7",
+             "#ead3c6",
+             "#f0b98d",
+             "#ef9708",
+             "#0fcfc0",
+             "#9cded6",
+             "#d5eae7",
+             "#f3e1eb",
+             "#f6c4e1",
+             "#f79cd4",
+             '#7f7f7f',
+             "#c7c7c7",
+             "#1CE6FF",
+             "#336600"])}
+    elif color_palette == "cell_type_20":
+        # https://github.com/vega/vega/wiki/Scales#scale-range-literals (some adjusted)
+        new_color_dict = {key: value for key, value in zip(
+            new_categories,
+            ['#1f77b4',
+             '#ff7f0e',
+             '#279e68',
+             '#d62728',
+             '#aa40fc',
+             '#8c564b',
+             '#e377c2',
+             '#b5bd61',
+             '#17becf',
+             '#aec7e8',
+             '#ffbb78',
+             '#98df8a',
+             '#ff9896',
+             '#c5b0d5',
+             '#c49c94',
+             '#f7b6d2',
+             '#dbdb8d',
+             '#9edae5',
+             '#ad494a',
+             '#8c6d31'])}
+    elif color_palette == "cell_type_10":
+        # scanpy vega10
+        new_color_dict = {key: value for key, value in zip(
+            new_categories,
+            ['#7f7f7f',
+             '#ff7f0e',
+             '#279e68',
+             '#e377c2',
+             '#17becf',
+             '#8c564b',
+             '#d62728',
+             '#1f77b4',
+             '#b5bd61',
+             '#aa40fc'])}
+    elif color_palette == "batch":
+        # sns.color_palette("colorblind").as_hex()
+        new_color_dict = {key: value for key, value in zip(
+            new_categories,
+            ['#0173b2', '#d55e00', '#ece133', '#ca9161', '#fbafe4',
+             '#949494', '#de8f05', '#029e73', '#cc78bc', '#56b4e9',
+             '#F0F8FF', '#FAEBD7', '#00FFFF', '#7FFFD4', '#F0FFFF',
+             '#F5F5DC', '#FFE4C4', '#000000', '#FFEBCD', '#0000FF',
+             '#8A2BE2', '#A52A2A', '#DEB887', '#5F9EA0', '#7FFF00',
+             '#D2691E', '#FF7F50', '#6495ED', '#FFF8DC', '#DC143C'])}
+    elif color_palette == "default":
+        new_color_dict = {key: value for key, value in zip(new_categories, list(default_color_dict.values())[skip_default_colors:])}
+    for key, val in overwrite_color_dict.items():
+        new_color_dict[key] = val
+    return new_color_dict
+
+
+def plot_non_zero_gene_count_means_dist(
+        adata: AnnData,
+        genes: list,
+        gene_label: str):
+    """
+    Plot distribution of non zero gene count means in the adata over all 
+    specified genes.
+    """
+    gene_counts = adata[
+        :, [gene for gene in adata.var_names if gene in genes]].layers["counts"]
+    nz_gene_means = np.mean(
+        np.ma.masked_equal(gene_counts.toarray(), 0), axis=0).data
+    
+    sns.kdeplot(nz_gene_means)
+    plt.title(f"{gene_label} Genes Average Non-Zero Gene Counts per Gene")
+    plt.xlabel("Average Non-zero Gene Counts")
+    plt.ylabel("Gene Density")
+    plt.show()
+
+
+def compute_communication_gp_network(
+    gp_list: list,
+    model: NicheCompass,
+    group_key: str="niche",
+    filter_key: Optional[str]=None,
+    filter_cat: Optional[str]=None,
+    n_neighbors: int=90):
+    """
+    Compute a network of category aggregated cell-pair communication strengths.
+    
+    First, compute cell-cell communication potential scores for each cell.
+    Then dot product them and take into account neighborhoods to compute
+    cell-pair communication strengths. Then, normalize cell-pair communication
+    strengths.
+    
+    Parameters
+    ----------
+    gp_list:
+        List of GPs for which the cell-pair communication strengths are computed.
+    model:
+        A trained NicheCompass model.
+    group_key:
+        Key in ´adata.obs´ where the groups are stored over which the cell-pair
+        communication strengths will be aggregated.
+    filter_key:
+        Key in ´adata.obs´ that contains the category for which the results are
+        filtered.
+    filter_cat:
+        Category for which the results are filtered.
+    n_neighbors:
+        Number of neighbors for the gp-specific neighborhood graph.
+
+    Returns
+    ----------
+    network_df:
+        A pandas dataframe with aggregated, normalized cell-pair communication strengths.
+    """
+    # Compute neighborhood graph
+    compute_knn = True
+    if 'spatial_cci' in model.adata.uns.keys():
+        if model.adata.uns['spatial_cci']['params']['n_neighbors'] == n_neighbors:
+            compute_knn = False
+    if compute_knn:
+        sc.pp.neighbors(model.adata,
+                        n_neighbors=n_neighbors,
+                        use_rep="spatial",
+                        key_added="spatial_cci")
+    
+    gp_network_dfs = []
+    gp_summary_df = model.get_gp_summary()
+    for gp in gp_list:
+        gp_idx = model.adata.uns[model.gp_names_key_].tolist().index(gp)
+        active_gp_idx = model.adata.uns[model.active_gp_names_key_].tolist().index(gp)
+        gp_scores = model.adata.obsm[model.latent_key_][:, active_gp_idx]
+        gp_targets_cats = model.adata.varm[model.gp_targets_categories_mask_key_][:, gp_idx]
+        gp_sources_cats = model.adata.varm[model.gp_sources_categories_mask_key_][:, gp_idx]
+        targets_cats_label_encoder = model.adata.uns[model.targets_categories_label_encoder_key_]
+        sources_cats_label_encoder = model.adata.uns[model.sources_categories_label_encoder_key_]
+
+        sources_cat_idx_dict = {}
+        for source_cat, source_cat_label in sources_cats_label_encoder.items():
+            sources_cat_idx_dict[source_cat] = np.where(gp_sources_cats == source_cat_label)[0]
+
+        targets_cat_idx_dict = {}
+        for target_cat, target_cat_label in targets_cats_label_encoder.items():
+            targets_cat_idx_dict[target_cat] = np.where(gp_targets_cats == target_cat_label)[0]
+
+        # Get indices of all source and target genes
+        source_genes_idx = np.array([], dtype=np.int64)
+        for key in sources_cat_idx_dict.keys():
+            source_genes_idx = np.append(source_genes_idx,
+                                         sources_cat_idx_dict[key])
+        target_genes_idx = np.array([], dtype=np.int64)
+        for key in targets_cat_idx_dict.keys():
+            target_genes_idx = np.append(target_genes_idx,
+                                         targets_cat_idx_dict[key])
+
+        # Compute cell-cell communication potential scores
+        gp_source_scores = np.zeros((len(model.adata.obs), len(source_genes_idx)))
+        gp_target_scores = np.zeros((len(model.adata.obs), len(target_genes_idx)))
+
+        for i, source_gene_idx in enumerate(source_genes_idx):
+            source_gene = model.adata.var_names[source_gene_idx]
+            gp_source_scores[:, i] = (
+                model.adata[:, model.adata.var_names.tolist().index(source_gene)].X.toarray().flatten() / model.adata[:, model.adata.var_names.tolist().index(source_gene)].X.toarray().flatten().max() *
+                gp_summary_df[gp_summary_df["gp_name"] == gp]["gp_source_genes_weights"].values[0][gp_summary_df[gp_summary_df["gp_name"] == gp]["gp_source_genes"].values[0].index(source_gene)] *
+                gp_scores)
+
+        for j, target_gene_idx in enumerate(target_genes_idx):
+            target_gene = model.adata.var_names[target_gene_idx]
+            gp_target_scores[:, j] = (
+                model.adata[:, model.adata.var_names.tolist().index(target_gene)].X.toarray().flatten() / model.adata[:, model.adata.var_names.tolist().index(target_gene)].X.toarray().flatten().max() *
+                gp_summary_df[gp_summary_df["gp_name"] == gp]["gp_target_genes_weights"].values[0][gp_summary_df[gp_summary_df["gp_name"] == gp]["gp_target_genes"].values[0].index(target_gene)] *
+                gp_scores)
+
+        agg_gp_source_score = gp_source_scores.mean(1).astype("float32")
+        agg_gp_target_score = gp_target_scores.mean(1).astype("float32")
+        agg_gp_source_score[agg_gp_source_score < 0] = 0.
+        agg_gp_target_score[agg_gp_target_score < 0] = 0.
+
+        model.adata.obs[f"{gp}_source_score"] = agg_gp_source_score
+        model.adata.obs[f"{gp}_target_score"] = agg_gp_target_score
+        
+        del(gp_target_scores)
+        del(gp_source_scores)
+
+        agg_gp_source_score = sp.csr_matrix(agg_gp_source_score)
+        agg_gp_target_score = sp.csr_matrix(agg_gp_target_score)
+
+        model.adata.obsp[f"{gp}_connectivities"] = (model.adata.obsp["spatial_cci_connectivities"] > 0).multiply(
+            agg_gp_source_score.T.dot(agg_gp_target_score))
+
+        # Aggregate gp connectivities for each group
+        gp_network_df_pivoted = aggregate_obsp_matrix_per_cell_type(
+            adata=model.adata,
+            obsp_key=f"{gp}_connectivities",
+            cell_type_key=group_key,
+            group_key=filter_key,
+            agg_rows=True)
+
+        if filter_key is not None:
+            gp_network_df_pivoted = gp_network_df_pivoted.loc[filter_cat, :]
+
+        gp_network_df = gp_network_df_pivoted.melt(var_name="source", value_name="gp_score", ignore_index=False).reset_index()
+        gp_network_df.columns = ["source", "target", "strength"]
+
+        gp_network_df = gp_network_df.sort_values("strength", ascending=False)
+
+        # Normalize strength
+        min_value = gp_network_df["strength"].min()
+        max_value = gp_network_df["strength"].max()
+        gp_network_df["strength_unscaled"] = gp_network_df["strength"]
+        gp_network_df["strength"] = (gp_network_df["strength"] - min_value) / (max_value - min_value)
+        gp_network_df["strength"] = np.round(gp_network_df["strength"], 2)
+        gp_network_df = gp_network_df[gp_network_df["strength"] > 0]
+
+        gp_network_df["edge_type"] = gp
+        gp_network_dfs.append(gp_network_df)
+
+    network_df = pd.concat(gp_network_dfs, ignore_index=True)
+    return network_df
+
+
+def visualize_communication_gp_network(
+    adata,
+    network_df,
+    cat_colors,
+    edge_type_colors: Optional[dict]=None,
+    edge_width_scale: int=20.0,
+    node_size: int=500,
+    fontsize: int=14,
+    figsize: Tuple[int, int]=(18, 16),
+    plot_legend: bool=True,
+    save: bool=False,
+    save_path: str="communication_gp_network.svg",
+    show: bool=True,
+    text_space: float=1.3,
+    connection_style="arc3, rad = 0.1",
+    cat_key: str="niche",
+    edge_attr: str="strength"):
+    """
+    Visualize a communication gp network.
+    """
+    # Assuming you have unique edge types in your 'edge_type' column
+    edge_types = np.unique(network_df['edge_type'])
+    
+    if edge_type_colors is None:
+        # Colorblindness adjusted vega_10
+        # See https://github.com/theislab/scanpy/issues/387
+        vega_10 = list(map(colors.to_hex, cm.tab10.colors))
+        vega_10_scanpy = vega_10.copy()
+        vega_10_scanpy[2] = "#279e68"  # green
+        vega_10_scanpy[4] = "#aa40fc"  # purple
+        vega_10_scanpy[8] = "#b5bd61"  # kakhi
+        edge_type_colors = vega_10_scanpy
+
+    # Create a dictionary that maps edge types to colors
+    edge_type_color_dict = {edge_type: color for edge_type, color in zip(edge_types, edge_type_colors)}
+
+    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=figsize)
+    ax.axis("off")
+    G = nx.from_pandas_edgelist(
+        network_df,
+        source="source",
+        target="target",
+        edge_attr=["edge_type", edge_attr],
+        create_using=nx.DiGraph(),
+    )
+    pos = nx.circular_layout(G)
+
+    nx.set_node_attributes(G, cat_colors, "color")
+    node_color = nx.get_node_attributes(G, "color")
+
+    description = nx.draw_networkx_labels(G, pos, font_size=fontsize)
+    n = adata.obs[cat_key].nunique()
+    node_list = sorted(G.nodes())
+    angle = []
+    angle_dict = {}
+    for i, node in zip(range(n), node_list):
+        theta = 2.0 * np.pi * i / n
+        angle.append((np.cos(theta), np.sin(theta)))
+        angle_dict[node] = theta
+    pos = {}
+    for node_i, node in enumerate(node_list):
+        pos[node] = angle[node_i]
+
+    r = fig.canvas.get_renderer()
+    trans = plt.gca().transData.inverted()
+    for node, t in description.items():
+        bb = t.get_window_extent(renderer=r)
+        bbdata = bb.transformed(trans)
+        radius = text_space + bbdata.width / 2.0
+        position = (radius * np.cos(angle_dict[node]), radius * np.sin(angle_dict[node]))
+        t.set_position(position)
+        t.set_rotation(angle_dict[node] * 360.0 / (2.0 * np.pi))
+        t.set_clip_on(False)
+
+    edgelist = [(u, v) for u, v, e in G.edges(data=True) if u != v]
+    edge_colors = [edge_type_color_dict[edge_data['edge_type']] for u, v, edge_data in G.edges(data=True) if u != v]
+    width = [e[edge_attr] * edge_width_scale for u, v, e in G.edges(data=True) if u != v]
+
+    h2 = nx.draw_networkx(
+        G,
+        pos,
+        with_labels=False,
+        node_size=node_size,
+        edgelist=edgelist,
+        width=width,
+        edge_vmin=0.0,
+        edge_vmax=1.0,
+        edge_color=edge_colors,  # Use the edge type colors here
+        arrows=True,
+        arrowstyle="-|>",
+        arrowsize=20,
+        vmin=0.0,
+        vmax=1.0,
+        cmap=plt.cm.binary,  # Use a colormap for node colors if needed
+        node_color=list(node_color.values()),
+        ax=ax,
+        connectionstyle=connection_style,
+    )
+
+    #https://stackoverflow.com/questions/19877666/add-legends-to-linecollection-plot - uses plotted data to define the color but here we already have colors defined, so just need a Line2D object.
+    def make_proxy(clr, mappable, **kwargs):
+        return Line2D([0, 1], [0, 1], color=clr, **kwargs)
+
+    # generate proxies with the above function
+    proxies = [make_proxy(clr, h2, lw=5) for clr in set(edge_colors)]
+    labels = [edge.split("_")[0] + " GP" for edge in edge_types[::-1]]
+
+    if plot_legend:
+        lgd = plt.legend(proxies, labels, loc="lower left")
+
+    edgelist = [(u, v) for u, v, e in G.edges(data=True) if ((u == v))] + [(u, v) for u, v, e in G.edges(data=True) if ((u != v))]
+    edge_colors = [edge_type_color_dict[edge_data['edge_type']] for u, v, edge_data in G.edges(data=True) if u == v]
+    width = [e[edge_attr] * edge_width_scale for u, v, e in G.edges(data=True) if u == v] + [0 for u, v, e in G.edges(data=True) if ((u != v))]
+    nx.draw_networkx_edges(
+        G,
+        pos,
+        node_size=node_size,
+        edgelist=edgelist, 
+        width=width,
+        edge_vmin=0.0,
+        edge_vmax=1.0,
+        edge_color=edge_colors,
+        arrows=False,
+        arrowstyle="-|>",
+        arrowsize=20,
+        ax=ax,
+        connectionstyle=connection_style)
+    plt.tight_layout()
+    if save:
+        plt.savefig(save_path)
+    if show:
+        plt.show()
+    plt.close(fig)
+    plt.ion()