local.py
High-level description
The local.py file in the cassiopeia/plotting module provides functions for plotting phylogenetic trees locally using either Matplotlib (for static plots) or Plotly (for dynamic plots). These functions offer an alternative to using iTOL, a cloud-based tree plotting service, for visualizing CassiopeiaTree objects.
Code Structure
The code defines several helper functions for placing tree elements and annotations on a coordinate system, creating colorstrips for categorical and continuous variables, and generating an indel heatmap. These helper functions are then used by the main plotting functions, plot_matplotlib and plot_plotly, to generate the final tree visualizations.
References
This file references symbols from the cassiopeia.data module, specifically the CassiopeiaTree class, and the cassiopeia.plotting.utilities module for tree placement and coordinate transformations.
Symbols
compute_colorstrip_size
Description
This function calculates the width and height of colorstrip boxes based on the node coordinates, anchor coordinates, and the desired location of the colorstrip relative to the anchor. The height is determined to ensure tight arrangement of boxes, while the width is set to 5% of the tree depth.
Inputs
| Name | Type | Description |
|---|---|---|
| node_coords | Dict[str, Tuple[float, float]] | A dictionary mapping node names to their (x, y) coordinates. |
| anchor_coords | Dict[str, Tuple[float, float]] | A dictionary mapping anchor node names to their (x, y) coordinates. |
| loc | Literal[“left”, “right”, “up”, “down”, “polar”] | The location of the colorstrip relative to the anchor. |
Outputs
| Name | Type | Description |
|---|---|---|
| width, height | Tuple[float, float] | A tuple containing the calculated width and height of the colorstrip boxes. |
Internal Logic
The function iterates through the node coordinates to determine the minimum and maximum positions and depths based on the specified location. The width is calculated as 5% of the difference between the maximum and minimum depths. The height is calculated by dividing the difference between the maximum and minimum positions by the number of anchor coordinates.
create_categorical_colorstrip
Description
This function creates a colorstrip for a categorical variable, assigning a unique color from a specified colormap to each unique value.
Inputs
| Name | Type | Description |
|---|---|---|
| values | Dict[str, str] | A dictionary mapping node names to their categorical values. |
| anchor_coords | Dict[str, Tuple[float, float]] | A dictionary mapping anchor node names to their (x, y) coordinates. |
| width | float | The width of the colorstrip. |
| height | float | The height of each box in the colorstrip. |
| spacing | float | The padding between consecutive colorstrip boxes. |
| loc | Literal[“left”, “right”, “up”, “down”, “polar”] | The location of the colorstrip relative to the anchor. |
| cmap | Union[str, mpl.colors.Colormap] | The colormap to use for assigning colors. Defaults to “tab10”. |
| value_mapping | Optional[Dict[str, int]] | An optional dictionary mapping string values to integer indices for color assignment. |
Outputs
| Name | Type | Description |
|---|---|---|
| colorstrip, next_anchor_coords | Tuple[Dict[str, Tuple[List[float], List[float], Tuple[float, float, float], str]], Dict[str, Tuple[float, float]]] | A tuple containing two dictionaries: the first maps node names to their colorstrip box coordinates, color, and text label; the second maps node names to the next set of anchor coordinates for placing additional colorstrips. |
Internal Logic
The function first determines the unique categorical values and assigns integer indices to them, either using the provided value_mapping or generating a pseudo-random mapping. It then uses the utilities.place_colorstrip function to calculate the box coordinates and next anchor coordinates. Finally, it iterates through the input values, assigns a color from the colormap based on the integer index, and populates the colorstrip dictionary with the box coordinates, color, and text label for each node.
create_continuous_colorstrip
Description
This function creates a colorstrip for a continuous variable, assigning colors from a specified colormap based on the normalized values within a specified range.
Inputs
| Name | Type | Description |
|---|---|---|
| values | Dict[str, float] | A dictionary mapping node names to their continuous values. |
| anchor_coords | Dict[str, Tuple[float, float]] | A dictionary mapping anchor node names to their (x, y) coordinates. |
| width | float | The width of the colorstrip. |
| height | float | The height of each box in the colorstrip. |
| spacing | float | The padding between consecutive colorstrip boxes. |
| loc | Literal[“left”, “right”, “up”, “down”, “polar”] | The location of the colorstrip relative to the anchor. |
| cmap | Union[str, mpl.colors.Colormap] | The colormap to use for assigning colors. Defaults to “viridis”. |
| vmin | Optional[float] | The minimum value for the color scale. |
| vmax | Optional[float] | The maximum value for the color scale. |
Outputs
| Name | Type | Description |
|---|---|---|
| colorstrip, next_anchor_coords | Tuple[Dict[str, Tuple[List[float], List[float], Tuple[float, float, float], str]], Dict[str, Tuple[float, float]]] | A tuple containing two dictionaries: the first maps node names to their colorstrip box coordinates, color, and text label; the second maps node names to the next set of anchor coordinates for placing additional colorstrips. |
Internal Logic
The function first determines the minimum and maximum values for the color scale, using the provided vmin and vmax or the minimum and maximum values from the input values. It then uses the utilities.place_colorstrip function to calculate the box coordinates and next anchor coordinates. Finally, it iterates through the input values, normalizes them to the range [0, 1], assigns a color from the colormap based on the normalized value, and populates the colorstrip dictionary with the box coordinates, color, and text label for each node.
create_indel_heatmap
Description
This function creates an indel heatmap, visualizing the indels present in an allele table using a color mapping.
Inputs
| Name | Type | Description |
|---|---|---|
| allele_table | pd.DataFrame | The allele table containing indel information. |
| anchor_coords | Dict[str, Tuple[float, float]] | A dictionary mapping anchor node names to their (x, y) coordinates. |
| width | float | The width of the heatmap. |
| height | float | The height of each box in the heatmap. |
| spacing | float | The padding between consecutive heatmap boxes. |
| loc | Literal[“left”, “right”, “up”, “down”, “polar”] | The location of the heatmap relative to the anchor. |
| indel_colors | Optional[pd.DataFrame] | An optional DataFrame mapping indels to colors. |
| indel_priors | Optional[pd.DataFrame] | An optional DataFrame containing prior probabilities for each indel. |
| random_state | Optional[np.random.RandomState] | An optional random state for reproducibility. |
Outputs
| Name | Type | Description |
|---|---|---|
| heatmap, anchor_coords | Tuple[List[Dict[str, Tuple[List[float], List[float], Tuple[float, float, float], str]]], Dict[str, Tuple[float, float]]] | A tuple containing two elements: the first is a list of colorstrips, where each colorstrip is a dictionary mapping node names to their box coordinates, color, and text label; the second is a dictionary mapping node names to the next set of anchor coordinates for placing additional colorstrips. |
Internal Logic
The function first prepares the allele table and indel colors using the utilities.prepare_alleletable function. It then iterates through each site in the allele table, calculates the box coordinates and next anchor coordinates using the utilities.place_colorstrip function, and assigns a color to each indel based on the provided indel_colors or a randomly generated color mapping. The resulting colorstrip for each site is appended to the heatmap list.
create_clade_colors
Description
This function assigns colors to nodes and branches based on clade membership, using a provided dictionary mapping clade root nodes to colors.
Inputs
| Name | Type | Description |
|---|---|---|
| tree | CassiopeiaTree | The CassiopeiaTree object. |
| clade_colors | Dict[str, Tuple[float, float, float]] | A dictionary mapping clade root nodes to RGB color tuples. |
Outputs
| Name | Type | Description |
|---|---|---|
| node_colors, branch_colors | Tuple[Dict[str, Tuple[float, float, float]], Dict[Tuple[str, str], Tuple[float, float, float]]] | A tuple containing two dictionaries: the first maps node names to their assigned colors; the second maps branch tuples (parent, child) to their assigned colors. |
Internal Logic
The function first identifies all descendants of each clade root node. It then iterates through the clade root nodes in descending order of their descendant set size, assigning the corresponding color to all nodes and branches within each clade.
place_tree_and_annotations
Description
This helper function places the tree and all requested annotations on a coordinate system, preparing the data for plotting.
Inputs
| Name | Type | Description |
|---|---|---|
| tree | CassiopeiaTree | The CassiopeiaTree object to plot. |
| depth_key | Optional[str] | The node attribute to use as the depth of the nodes. |
| meta_data | Optional[List[str]] | A list of metadata columns to plot alongside the tree. |
| allele_table | Optional[pd.DataFrame] | The allele table to plot alongside the tree. |
| indel_colors | Optional[pd.DataFrame] | A DataFrame mapping indels to colors for the indel heatmap. |
| indel_priors | Optional[pd.DataFrame] | A DataFrame containing prior probabilities for each indel. |
| orient | Union[Literal[“up”, “down”, “left”, “right”], float] | The orientation of the tree. |
| extend_branches | bool | Whether to extend branch lengths to equalize distances from the root. |
| angled_branches | bool | Whether to display branches as angled lines. |
| add_root | bool | Whether to add a synthetic root node. |
| colorstrip_width | Optional[float] | The width of the colorstrips. |
| colorstrip_spacing | Optional[float] | The spacing between consecutive colorstrips. |
| clade_colors | Optional[Dict[str, Tuple[float, float, float]]] | A dictionary mapping clade root nodes to colors. |
| continuous_cmap | Union[str, mpl.colors.Colormap] | The colormap to use for continuous variables. |
| vmin | Optional[float] | The minimum value for the continuous color scale. |
| vmax | Optional[float] | The maximum value for the continuous color scale. |
| categorical_cmap | Union[str, mpl.colors.Colormap] | The colormap to use for categorical variables. |
| value_mapping | Optional[Dict[str, int]] | A dictionary mapping string values to integer indices for categorical color assignment. |
| random_state | Optional[np.random.RandomState] | A random state for reproducibility. |
Outputs
| Name | Type | Description |
|---|---|---|
| node_coords, branch_coords, node_colors, branch_colors, colorstrips | Tuple[Dict, Dict, Dict, Dict, List] | A tuple containing five elements: node coordinates, branch coordinates, node colors, branch colors, and a list of colorstrips. |
Internal Logic
The function first places the tree on the coordinate system using the utilities.place_tree function. It then calculates the initial anchor coordinates for the colorstrips, which are the coordinates of the leaf nodes. Next, it creates the indel heatmap, if an allele table is provided, and iterates through the metadata columns, creating colorstrips for each column based on its data type. Finally, it assigns colors to nodes and branches based on clade membership, if clade_colors is provided.
plot_matplotlib
Description
This function generates a static plot of a CassiopeiaTree using Matplotlib.
Inputs
| Name | Type | Description |
|---|---|---|
| tree | CassiopeiaTree | The CassiopeiaTree object to plot. |
| depth_key | Optional[str] | The node attribute to use as the depth of the nodes. |