visualization.py

Go to the documentation of this file.

import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
import plotly.graph_objs as go
from typing import Dict, Any, Optional
 
def drawCellLineage(data: 'Dict[str, Any]', cell_attribute: str, use_plotly: bool = False, color_scale: str = 'Spectral_r', discrete_colors: Optional[Dict[Any, str]] = None) -> None:
    """
    Draws the lineage of cells in the simulation based on a specified cell attribute.
 
    Parameters
    ----------
    - data : dict
        The dictionary containing cell lineage data and enumeration map. It should have the keys 'data' and 'enumMap'.
    - cell_attribute : str
        The name of the cell attribute to visualize. This can be a discrete or continuous attribute.
    - use_plotly : bool, optional
        If True, use Plotly for visualization. If False, use Matplotlib. Default is False.
    - color_scale : str, optional
        The color scale to use for visualization. Default is 'Spectral_r'.
    - discrete_colors : dict, optional
        A dictionary mapping discrete attribute values to specific colors. If provided, this will be used for coloring discrete attributes.
 
    Returns
    -------
    - None
        This function does not return any value. It generates and displays a lineage plot.
 
    Example
    -------
    ```python
    simu = Simu(params)
    lineage_data = simu.recordLineage()
    discrete_colors = {'A': '#1f77b4', 'B': '#ff7f0e', 'C': '#2ca02c'}
    drawCellLineage(lineage_data, 'state', use_plotly=True, color_scale='Viridis', discrete_colors=discrete_colors)
    ```
 
    Notes
    -----
    - If the `cell_attribute` is an Enum, the corresponding values are replaced and the data type is set to 'category'.
    - For continuous attributes, a color map is used to represent the range of values.
    """
    charac_mapping = None
    cell_attribute=cell_attribute.lower()
    if cell_attribute in data['enumMap']:
        charac_mapping = data['enumMap'][cell_attribute]
    data = data['data']
    # Initialiser le graphe
    G = nx.DiGraph()
 
    # Ajouter un nœud fantôme
    G.add_node(-1, birth_step=-1, death_step=-1, characteristic=[])
 
    # Ajouter des nœuds et des arêtes
    for cell in data:
        G.add_node(cell['id'], birth_step=cell['birth_step'], death_step=cell['death_step'], characteristic=cell[cell_attribute])
        if cell['mother_id'] != -1:
            G.add_edge(cell['mother_id'], cell['id'])
        else:
            G.add_edge(-1, cell['id'])
 
    # Fonction pour identifier les arbres et calculer leurs hauteurs
    def get_tree_positions(G, y_spacing=1):
        pos = {}
        widths = {}
        y_offset = 0
 
        def calculate_widths(node):
            if node not in widths:
                if G.out_degree(node) == 0:
                    widths[node] = 1
                else:
                    widths[node] = sum(calculate_widths(child) for child in G.successors(node))
            return widths[node]
 
        def assign_position(node, x, y):
            pos[node] = (x, y)
            return y + y_spacing
 
        def dfs(node, x, y):
            if node in pos:
                return y
            width = widths[node]
            start_y = y
            for child in G.successors(node):
                child_width = widths[child]
                child_x = G.nodes[child]['birth_step']
                y = dfs(child, child_x, start_y)
                start_y += child_width * y_spacing
            if G.out_degree(node) > 0:
                y = (pos[list(G.successors(node))[0]][1] + pos[list(G.successors(node))[-1]][1]) / 2
            y = assign_position(node, x, y)
            return y
 
        # Calculer les largeurs des sous-arbres
        for node in G.nodes():
            calculate_widths(node)
 
        # Positionner les nœuds sans parents en premier
        y_offset = dfs(-1, -1, y_offset) + y_spacing
 
        return pos
 
    # Utiliser un espacement vertical plus grand pour éviter les superpositions
    pos = get_tree_positions(G, y_spacing=2)
 
    # Calculer le nombre de feuilles
    num_leaves = sum(1 for node in G.nodes if G.out_degree(node) == 0 and node != -1)
 
    # Déterminer la hauteur du plot en fonction du nombre de feuilles
    plot_height = num_leaves * 30  # Ajustez le facteur de multiplication selon vos besoins
 
    # Extraire les positions pour les nœuds
    node_x = []
    node_y = []
    vertical_lines = []
    horizontal_lines = []
    life_traces = []
 
    # Déterminer l'échelle de couleurs pour les caractéristiques discrètes
    if charac_mapping is not None:
        if discrete_colors:
            color_mapping = {v:discrete_colors[k] for v,k in charac_mapping.items()}
            norm = mcolors.BoundaryNorm(boundaries=np.arange(len(color_mapping) + 1) - 0.5, ncolors=len(color_mapping))
        else:
            unique_chars = sorted(set(val for cell in data for val in cell[cell_attribute]))
            colors = plt.cm.tab10(np.linspace(0, 1, len(unique_chars)))
            color_mapping = {val: mcolors.rgb2hex(colors[i]) for i, val in enumerate(unique_chars)}
            norm = mcolors.BoundaryNorm(boundaries=np.arange(len(unique_chars) + 1) - 0.5, ncolors=len(unique_chars))
    else:
        cmap = plt.get_cmap(color_scale)
        vmin = np.inf
        vmax = -np.inf
        for cell in data:
            vmin = min(np.min(cell[cell_attribute]), vmin)
            vmax = max(np.max(cell[cell_attribute]), vmax)
        norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
 
    for node in G.nodes():
        if node == -1:
            continue
        x, y = pos[node]
        node_x.append(x)
        node_y.append(y)
        birth_step = G.nodes[node]['birth_step']
        death_step = G.nodes[node]['death_step']
 
        characteristic = G.nodes[node]['characteristic']
        steps = list(range(birth_step, death_step))
 
        # Fractionner les lignes seulement si la caractéristique change
        if charac_mapping is not None:
            current_char = characteristic[0]
            current_color = color_mapping[current_char]
        else:
            current_color = mcolors.rgb2hex(cmap(norm(characteristic[0])))
        start_step = birth_step
        for i in range(1, len(steps)):
            if charac_mapping is not None:
                new_char = characteristic[i]
                new_color = color_mapping[new_char]
                if new_char != current_char:
                    vertical_lines.append(((start_step, y), (steps[i], y), current_color))
                    current_char = new_char
                    current_color = new_color
                    start_step = steps[i]
            else:
                new_color = mcolors.rgb2hex(cmap(norm(characteristic[i])))
                if new_color != current_color:
                    vertical_lines.append(((start_step, y), (steps[i], y), current_color))
                    current_color = new_color
                    start_step = steps[i]
        vertical_lines.append(((start_step, y), (death_step, y), current_color))
 
    for edge in G.edges():
        parent = edge[0]
        child = edge[1]
        if parent == -1:
            continue
        parent_x, parent_y = pos[parent]
        child_x, child_y = pos[child]
        division_step = G.nodes[child]['birth_step']
        horizontal_lines.append(((division_step, parent_y), (division_step, child_y)))
 
    if use_plotly:
        # Créer des traces pour les lignes de vie (lignes verticales) en premier
        for line in vertical_lines:
            (x0, y0), (x1, y1), color = line
            life_traces.append(go.Scatter(
                x=[x0, x1],
                y=[y0, y0],
                mode='lines',
                line=dict(width=4, color=color),
                hoverinfo='none',
                showlegend=False
            ))
 
        # Créer des traces pour les arêtes (lignes horizontales pour les divisions)
        edge_x = []
        edge_y = []
 
        for line in horizontal_lines:
            (x0, y0), (x1, y1) = line
            edge_x.extend([x0, x1, None])
            edge_y.extend([y0, y1, None])
 
        edge_trace = go.Scatter(
            x=edge_x, y=edge_y,
            line=dict(width=2, color='#888'),
            hoverinfo='none',
            showlegend=False,
            mode='lines'
        )
 
        # Ajouter des traces pour les légendes des caractéristiques discrètes
        if charac_mapping is not None:
            for char_val, char_name in charac_mapping.items():
                life_traces.append(go.Scatter(
                    x=[None],
                    y=[None],
                    mode='lines',
                    line=dict(width=2, color=color_mapping[char_val]),
                    hoverinfo='none',
                    showlegend=True,
                    name=char_name
                ))
        else:
            # Ajouter une colorbar pour les lignes verticales
            colorbar_trace = go.Scatter(
                x=[None],
                y=[None],
                mode='markers',
                marker=dict(
                    colorscale=color_scale,
                    cmin=vmin,
                    cmax=vmax,
                    colorbar=dict(
                        thickness=15,
                        title=cell_attribute,
                        xanchor='left',
                        titleside='right'
                    )
                ),
                showlegend=False
            )
            life_traces.append(colorbar_trace)
 
        # Créer la figure finale
        fig = go.Figure(data=life_traces + [edge_trace],
                        layout=go.Layout(
                            title='Lineage tree of Cells',
                            titlefont_size=16,
                            showlegend=True,
                            hovermode='closest',
                            margin=dict(b=20, l=5, r=5, t=40),
                            xaxis=dict(showgrid=True, zeroline=True, showticklabels=True, range=[min(node_x), max(node_x)]),
                            yaxis=dict(showgrid=True, zeroline=True, showticklabels=False, range=[min(node_y) - 1, max(node_y) + 1]),
                            height=plot_height  # Ajuster la hauteur du plot
                        )
                        )
 
        fig.show()
    else:
        fig, ax = plt.subplots(figsize=(10, plot_height / 100))  # Ajuster la hauteur du plot
 
        # Tracer les lignes de vie (lignes verticales) en premier
        for line in vertical_lines:
            (x0, y0), (x1, y1), color = line
            ax.plot([x0, x1], [y0, y0], color=color, linewidth=4, zorder=3)
 
        # Tracer les lignes horizontales
        for line in horizontal_lines:
            (x0, y0), (x1, y1) = line
            ax.plot([x0, x1], [y0, y1], color='#888', linewidth=2, zorder=1)
 
        # Tracer les nœuds
        # ax.scatter(node_x, node_y, color='black', s=10, zorder=2)
 
        if charac_mapping is not None:
            # Ajouter des légendes pour les caractéristiques discrètes
            handles = [plt.Line2D([0], [0], color=color_mapping[char_val], lw=2, label=char_name)
                       for char_val, char_name in charac_mapping.items()]
            ax.legend(handles=handles, title=cell_attribute, loc='center left', bbox_to_anchor=(1, 0.5))
        else:
            # Ajouter une colorbar
            sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
            sm.set_array([])
            cbar = plt.colorbar(sm, ax=ax, fraction=0.02, pad=0.04)  # Ajuster fraction et pad
            cbar.set_label(cell_attribute)
 
        ax.set_title('Lineage tree of Cells')
        ax.set_xlabel('Step')
        ax.set_ylabel('Cells')
        ax.grid(True)
 
        # Ajuster les limites des axes pour s'adapter parfaitement au contenu
        ax.set_xlim(min(node_x), max(node_x))
        ax.set_ylim(min(node_y) - 1, max(node_y) + 1)
 
        plt.tight_layout()
        plt.show()