> ## Documentation Index
> Fetch the complete documentation index at: https://demo.agenticlabs.com/llms.txt
> Use this file to discover all available pages before exploring further.

# 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. |
