phylogenie 3.1.5__tar.gz → 3.1.9__tar.gz

{phylogenie-3.1.5/src/phylogenie.egg-info → phylogenie-3.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: phylogenie
-Version: 3.1.5
+Version: 3.1.9
 Summary: Generate phylogenetic datasets with minimal setup effort
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{phylogenie-3.1.5 → phylogenie-3.1.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "phylogenie"
-version = "3.1.5"
+version = "3.1.9"
 description = "Generate phylogenetic datasets with minimal setup effort"
 readme = "README.md"
 requires-python = ">=3.10"

{phylogenie-3.1.5 → phylogenie-3.1.9}/src/phylogenie/__init__.py

@@ -1,4 +1,11 @@
-from phylogenie.draw import Coloring, draw_tree
+from phylogenie.draw import (
+    draw_colored_dated_tree_categorical,
+    draw_colored_dated_tree_continuous,
+    draw_colored_tree_categorical,
+    draw_colored_tree_continuous,
+    draw_dated_tree,
+    draw_tree,
+)
 from phylogenie.generators import (
     AliSimDatasetGenerator,
     BDEITreeDatasetGenerator,
@@ -11,8 +18,8 @@ from phylogenie.generators import (
     FBDTreeDatasetGenerator,
     TreeDatasetGeneratorConfig,
 )
-from phylogenie.io import load_fasta
-from phylogenie.msa import MSA
+from phylogenie.io import dump_fasta, load_fasta
+from phylogenie.msa import MSA, Sequence
 from phylogenie.skyline import (
     SkylineMatrix,
     SkylineMatrixCoercible,
@@ -64,7 +71,11 @@ from phylogenie.treesimulator import (
 )
 
 __all__ = [
-    "Coloring",
+    "draw_colored_dated_tree_categorical",
+    "draw_colored_dated_tree_continuous",
+    "draw_colored_tree_categorical",
+    "draw_colored_tree_continuous",
+    "draw_dated_tree",
     "draw_tree",
     "AliSimDatasetGenerator",
     "BDEITreeDatasetGenerator",
@@ -76,8 +87,10 @@ __all__ = [
     "EpidemiologicalTreeDatasetGenerator",
     "FBDTreeDatasetGenerator",
     "TreeDatasetGeneratorConfig",
+    "dump_fasta",
     "load_fasta",
     "MSA",
+    "Sequence",
     "SkylineMatrix",
     "SkylineMatrixCoercible",
     "SkylineParameter",

phylogenie-3.1.9/src/phylogenie/draw.py

@@ -0,0 +1,691 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Literal, overload
+
+import matplotlib.colors as mcolors
+import matplotlib.dates as mdates
+import matplotlib.patches as mpatches
+import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
+from matplotlib.colors import Colormap
+from mpl_toolkits.axes_grid1.inset_locator import inset_axes  # pyright: ignore
+
+from phylogenie.treesimulator import (
+    Tree,
+    get_node_ages,
+    get_node_depth_levels,
+    get_node_depths,
+)
+
+
+@dataclass
+class CalibrationNode:
+    node: Tree
+    date: datetime
+
+
+Color = str | tuple[float, float, float] | tuple[float, float, float, float]
+
+
+def draw_tree(
+    tree: Tree,
+    ax: Axes | None = None,
+    colors: Color | dict[Tree, Color] = "black",
+    backward_time: bool = False,
+) -> Axes:
+    """
+    Draw a phylogenetic tree with colored branches.
+
+    Parameters
+    ----------
+    tree : Tree
+        The phylogenetic tree to draw.
+    ax : Axes | None, optional
+        The matplotlib Axes to draw on. If None, uses the current Axes.
+    colors : Color | dict[Tree, Color], optional
+        A single color for all branches or a dictionary mapping each node to a color.
+    backward_time : bool, optional
+        If True, the x-axis is inverted to represent time going backward.
+
+    Returns
+    -------
+    Axes
+        The Axes with the drawn tree.
+    """
+    if ax is None:
+        ax = plt.gca()
+
+    if not isinstance(colors, dict):
+        colors = {node: colors for node in tree}
+
+    xs = (
+        get_node_ages(tree)
+        if backward_time
+        else get_node_depth_levels(tree)
+        if any(node.branch_length is None for node in tree.iter_descendants())
+        else get_node_depths(tree)
+    )
+
+    ys: dict[Tree, float] = {node: i for i, node in enumerate(tree.get_leaves())}
+    for node in tree.postorder_traversal():
+        if node.is_internal():
+            ys[node] = sum(ys[child] for child in node.children) / len(node.children)
+
+    if tree.branch_length is not None:
+        xmin = xs[tree] + tree.branch_length if backward_time else 0
+        ax.hlines(y=ys[tree], xmin=xmin, xmax=xs[tree], color=colors[tree])  # pyright: ignore
+    for node in tree:
+        x1, y1 = xs[node], ys[node]
+        for child in node.children:
+            x2, y2 = xs[child], ys[child]
+            ax.hlines(y=y2, xmin=x1, xmax=x2, color=colors[child])  # pyright: ignore
+            ax.vlines(x=x1, ymin=y1, ymax=y2, color=colors[child])  # pyright: ignore
+
+    if backward_time:
+        ax.invert_xaxis()
+
+    ax.set_yticks([])  # pyright: ignore
+    return ax
+
+
+def _depth_to_date(
+    depth: float, calibration_nodes: tuple[CalibrationNode, CalibrationNode]
+) -> datetime:
+    """
+    Convert a depth value to a date using linear interpolation between two calibration nodes.
+
+    Parameters
+    ----------
+    depth : float
+        The depth value to convert.
+    calibration_nodes : tuple[CalibrationNode, CalibrationNode]
+        Two calibration nodes defining the mapping from depth to date.
+
+    Returns
+    -------
+    datetime
+        The interpolated date corresponding to the given depth.
+    """
+    node1, node2 = calibration_nodes
+    depth1, depth2 = node1.node.depth, node2.node.depth
+    date1, date2 = node1.date, node2.date
+    return date1 + (depth - depth1) * (date2 - date1) / (depth2 - depth1)
+
+
+def draw_dated_tree(
+    tree: Tree,
+    calibration_nodes: tuple[CalibrationNode, CalibrationNode],
+    ax: Axes | None = None,
+    colors: Color | dict[Tree, Color] = "black",
+) -> Axes:
+    """
+    Draw a phylogenetic tree with branches positioned according to calibrated dates.
+
+    Parameters
+    ----------
+    tree : Tree
+        The phylogenetic tree to draw.
+    calibration_nodes : tuple[CalibrationNode, CalibrationNode]
+        Two calibration nodes defining the mapping from depth to date.
+    ax : Axes | None, optional
+        The matplotlib Axes to draw on. If None, uses the current Axes.
+    colors : Color | dict[Tree, Color], optional
+        A single color for all branches or a dictionary mapping each node to a color.
+
+    Returns
+    -------
+    Axes
+        The Axes with the drawn dated tree.
+    """
+    if ax is None:
+        ax = plt.gca()
+
+    if not isinstance(colors, dict):
+        colors = {node: colors for node in tree}
+
+    xs = {
+        node: _depth_to_date(depth=depth, calibration_nodes=calibration_nodes)
+        for node, depth in get_node_depths(tree).items()
+    }
+
+    ys: dict[Tree, float] = {node: i for i, node in enumerate(tree.get_leaves())}
+    for node in tree.postorder_traversal():
+        if node.is_internal():
+            ys[node] = sum(ys[child] for child in node.children) / len(node.children)
+
+    if tree.branch_length is not None:
+        origin_date = _depth_to_date(depth=0, calibration_nodes=calibration_nodes)
+        ax.hlines(  # pyright: ignore
+            y=ys[tree],
+            xmin=mdates.date2num(origin_date),  # pyright: ignore
+            xmax=mdates.date2num(xs[tree]),  # pyright: ignore
+            color=colors[tree],
+        )
+    for node in tree:
+        x1, y1 = xs[node], ys[node]
+        for child in node.children:
+            x2, y2 = xs[child], ys[child]
+            ax.hlines(  # pyright: ignore
+                y=y2,
+                xmin=mdates.date2num(x1),  # pyright: ignore
+                xmax=mdates.date2num(x2),  # pyright: ignore
+                color=colors[child],
+            )
+            ax.vlines(x=mdates.date2num(x1), ymin=y1, ymax=y2, color=colors[child])  # pyright: ignore
+
+    ax.xaxis.set_major_locator(mdates.AutoDateLocator())
+    ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m-%d"))
+    ax.tick_params(axis="x", labelrotation=45)  # pyright: ignore
+
+    ax.set_yticks([])  # pyright: ignore
+    return ax
+
+
+def _init_colored_tree_categorical(
+    tree: Tree,
+    color_by: str,
+    ax: Axes | None = None,
+    default_color: Color = "black",
+    colormap: str | Colormap = "tab20",
+    show_legend: bool = True,
+    labels: dict[Any, str] | None = None,
+    legend_kwargs: dict[str, Any] | None = None,
+) -> tuple[Axes, dict[Tree, Color]]:
+    """
+    Initialize colors for drawing a tree based on categorical metadata.
+
+    Parameters
+    ----------
+    tree : Tree
+        The phylogenetic tree.
+    color_by : str
+        The metadata key to color branches by.
+    ax : Axes | None, optional
+        The matplotlib Axes to draw on. If None, uses the current Axes.
+    default_color : Color, optional
+        The color to use for nodes without the specified metadata.
+    colormap : str | Colormap, optional
+        The colormap to use for coloring categories. Defaults to 'tab20'.
+    show_legend : bool, optional
+        Whether to display a legend for the categories.
+    labels : dict[Any, str] | None, optional
+        A mapping from category values to labels for the legend.
+    legend_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the legend.
+
+    Returns
+    -------
+    tuple[Axes, dict[Tree, Color]]
+        The Axes and a dictionary mapping each node to its assigned color.
+    """
+    if ax is None:
+        ax = plt.gca()
+
+    if isinstance(colormap, str):
+        colormap = plt.get_cmap(colormap)
+
+    features = {node: node[color_by] for node in tree if color_by in node.metadata}
+    feature_colors = {
+        f: mcolors.to_hex(colormap(i)) for i, f in enumerate(set(features.values()))
+    }
+    colors = {
+        node: feature_colors[features[node]] if node in features else default_color
+        for node in tree
+    }
+
+    if show_legend:
+        legend_handles = [
+            mpatches.Patch(
+                color=feature_colors[f],
+                label=str(f) if labels is None else labels[f],
+            )
+            for f in feature_colors
+        ]
+        if any(color_by not in node.metadata for node in tree):
+            legend_handles.append(mpatches.Patch(color=default_color, label="NA"))
+        if legend_kwargs is None:
+            legend_kwargs = {}
+        ax.legend(handles=legend_handles, **legend_kwargs)  # pyright: ignore
+
+    return ax, colors
+
+
+def draw_colored_tree_categorical(
+    tree: Tree,
+    color_by: str,
+    ax: Axes | None = None,
+    backward_time: bool = False,
+    default_color: Color = "black",
+    colormap: str | Colormap = "tab20",
+    show_legend: bool = True,
+    labels: dict[Any, str] | None = None,
+    legend_kwargs: dict[str, Any] | None = None,
+):
+    """
+    Draw a phylogenetic tree with branches colored based on categorical metadata.
+
+    Parameters
+    ----------
+    tree : Tree
+        The phylogenetic tree to draw.
+    color_by : str
+        The metadata key to color branches by.
+    ax : Axes | None, optional
+        The matplotlib Axes to draw on. If None, uses the current Axes.
+    backward_time : bool, optional
+        If True, the x-axis is inverted to represent time going backward.
+    default_color : Color, optional
+        The color to use for nodes without the specified metadata.
+    colormap : str | Colormap, optional
+        The colormap to use for coloring categories. Defaults to 'tab20'.
+    show_legend : bool, optional
+        Whether to display a legend for the categories.
+    labels : dict[Any, str] | None, optional
+        A mapping from category values to labels for the legend.
+    legend_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the legend.
+
+    Returns
+    -------
+    Axes
+        The Axes with the drawn colored tree.
+    """
+    ax, colors = _init_colored_tree_categorical(
+        tree=tree,
+        color_by=color_by,
+        ax=ax,
+        default_color=default_color,
+        colormap=colormap,
+        show_legend=show_legend,
+        labels=labels,
+        legend_kwargs=legend_kwargs,
+    )
+    return draw_tree(tree=tree, ax=ax, colors=colors, backward_time=backward_time)
+
+
+def draw_colored_dated_tree_categorical(
+    tree: Tree,
+    calibration_nodes: tuple[CalibrationNode, CalibrationNode],
+    color_by: str,
+    ax: Axes | None = None,
+    default_color: Color = "black",
+    colormap: str | Colormap = "tab20",
+    show_legend: bool = True,
+    labels: dict[Any, str] | None = None,
+    legend_kwargs: dict[str, Any] | None = None,
+) -> Axes:
+    """
+    Draw a dated phylogenetic tree with branches colored based on categorical metadata.
+
+    Parameters
+    ----------
+    tree : Tree
+        The phylogenetic tree to draw.
+    calibration_nodes : tuple[CalibrationNode, CalibrationNode]
+        Two calibration nodes defining the mapping from depth to date.
+    color_by : str
+        The metadata key to color branches by.
+    ax : Axes | None, optional
+        The matplotlib Axes to draw on. If None, uses the current Axes.
+    default_color : Color, optional
+        The color to use for nodes without the specified metadata.
+    colormap : str | Colormap, optional
+        The colormap to use for coloring categories. Defaults to 'tab20'.
+    show_legend : bool, optional
+        Whether to display a legend for the categories.
+    labels : dict[Any, str] | None, optional
+        A mapping from category values to labels for the legend.
+    legend_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the legend.
+
+    Returns
+    -------
+    Axes
+        The Axes with the drawn colored dated tree.
+    """
+    ax, colors = _init_colored_tree_categorical(
+        tree=tree,
+        color_by=color_by,
+        ax=ax,
+        default_color=default_color,
+        colormap=colormap,
+        show_legend=show_legend,
+        labels=labels,
+        legend_kwargs=legend_kwargs,
+    )
+    return draw_dated_tree(
+        tree=tree, calibration_nodes=calibration_nodes, ax=ax, colors=colors
+    )
+
+
+@overload
+def _init_colored_tree_continuous(
+    tree: Tree,
+    color_by: str,
+    ax: Axes | None = ...,
+    default_color: Color = ...,
+    colormap: str | Colormap = ...,
+    vmin: float | None = ...,
+    vmax: float | None = ...,
+    *,
+    show_hist: Literal[False],
+    hist_kwargs: dict[str, Any] | None = ...,
+    hist_axes_kwargs: dict[str, Any] | None = ...,
+) -> tuple[Axes, dict[Tree, Color]]: ...
+@overload
+def _init_colored_tree_continuous(
+    tree: Tree,
+    color_by: str,
+    ax: Axes | None = ...,
+    default_color: Color = ...,
+    colormap: str | Colormap = ...,
+    vmin: float | None = ...,
+    vmax: float | None = ...,
+    *,
+    show_hist: Literal[True] = True,
+    hist_kwargs: dict[str, Any] | None = ...,
+    hist_axes_kwargs: dict[str, Any] | None = ...,
+) -> tuple[Axes, dict[Tree, Color], Axes]: ...
+def _init_colored_tree_continuous(
+    tree: Tree,
+    color_by: str,
+    ax: Axes | None = None,
+    default_color: Color = "black",
+    colormap: str | Colormap = "viridis",
+    vmin: float | None = None,
+    vmax: float | None = None,
+    show_hist: bool = True,
+    hist_kwargs: dict[str, Any] | None = None,
+    hist_axes_kwargs: dict[str, Any] | None = None,
+) -> tuple[Axes, dict[Tree, Color]] | tuple[Axes, dict[Tree, Color], Axes]:
+    """
+    Initialize colors for drawing a tree based on continuous metadata.
+
+    Parameters
+    ----------
+    tree : Tree
+        The phylogenetic tree.
+    color_by : str
+        The metadata key to color branches by.
+    ax : Axes | None, optional
+        The matplotlib Axes to draw on. If None, uses the current Axes.
+    default_color : Color, optional
+        The color to use for nodes without the specified metadata.
+    colormap : str | Colormap, optional
+        The colormap to use for coloring continuous values. Defaults to 'viridis'.
+    vmin : float | None, optional
+        The minimum value for normalization. If None, uses the minimum of the data.
+    vmax : float | None, optional
+        The maximum value for normalization. If None, uses the maximum of the data.
+    show_hist : bool, optional
+        Whether to display a histogram of the continuous values.
+    hist_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the histogram.
+    hist_axes_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to define the histogram Axes.
+
+    Returns
+    -------
+    tuple[Axes, dict[Tree, Color]] | tuple[Axes, dict[Tree, Color], Axes]
+        The Axes, a dictionary mapping each node to its assigned color,
+        and optionally the histogram Axes if `show_hist` is True.
+    """
+    if ax is None:
+        ax = plt.gca()
+
+    if isinstance(colormap, str):
+        colormap = plt.get_cmap(colormap)
+
+    features = {node: node[color_by] for node in tree if color_by in node.metadata}
+    values = list(features.values())
+    vmin = min(values) if vmin is None else vmin
+    vmax = max(values) if vmax is None else vmax
+    norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
+    colors = {
+        node: colormap(norm(float(features[node])))
+        if node in features
+        else default_color
+        for node in tree
+    }
+
+    if show_hist:
+        default_hist_axes_kwargs = {"width": "25%", "height": "25%"}
+        if hist_axes_kwargs is not None:
+            default_hist_axes_kwargs.update(hist_axes_kwargs)
+        hist_ax = inset_axes(ax, **default_hist_axes_kwargs)  # pyright: ignore
+
+        hist_kwargs = {} if hist_kwargs is None else hist_kwargs
+        _, bins, patches = hist_ax.hist(values, **hist_kwargs)  # pyright: ignore
+
+        for patch, b0, b1 in zip(patches, bins[:-1], bins[1:]):  # pyright: ignore
+            midpoint = (b0 + b1) / 2  # pyright: ignore
+            patch.set_facecolor(colormap(norm(midpoint)))  # pyright: ignore
+        return ax, colors, hist_ax  # pyright: ignore
+
+    sm = plt.cm.ScalarMappable(cmap=colormap, norm=norm)
+    ax.get_figure().colorbar(sm, ax=ax)  # pyright: ignore
+    return ax, colors
+
+
+@overload
+def draw_colored_tree_continuous(
+    tree: Tree,
+    color_by: str,
+    ax: Axes | None = None,
+    backward_time: bool = False,
+    default_color: Color = "black",
+    colormap: str | Colormap = "viridis",
+    vmin: float | None = None,
+    vmax: float | None = None,
+    *,
+    show_hist: Literal[False],
+    hist_kwargs: dict[str, Any] | None = None,
+    hist_axes_kwargs: dict[str, Any] | None = None,
+) -> Axes: ...
+@overload
+def draw_colored_tree_continuous(
+    tree: Tree,
+    color_by: str,
+    ax: Axes | None = None,
+    backward_time: bool = False,
+    default_color: Color = "black",
+    colormap: str | Colormap = "viridis",
+    vmin: float | None = None,
+    vmax: float | None = None,
+    show_hist: Literal[True] = True,
+    hist_kwargs: dict[str, Any] | None = None,
+    hist_axes_kwargs: dict[str, Any] | None = None,
+) -> tuple[Axes, Axes]: ...
+def draw_colored_tree_continuous(
+    tree: Tree,
+    color_by: str,
+    ax: Axes | None = None,
+    backward_time: bool = False,
+    default_color: Color = "black",
+    colormap: str | Colormap = "viridis",
+    vmin: float | None = None,
+    vmax: float | None = None,
+    show_hist: bool = True,
+    hist_kwargs: dict[str, Any] | None = None,
+    hist_axes_kwargs: dict[str, Any] | None = None,
+) -> Axes | tuple[Axes, Axes]:
+    """
+    Draw a phylogenetic tree with branches colored based on continuous metadata.
+
+    Parameters
+    ----------
+    tree : Tree
+        The phylogenetic tree to draw.
+    color_by : str
+        The metadata key to color branches by.
+    ax : Axes | None, optional
+        The matplotlib Axes to draw on. If None, uses the current Axes.
+    backward_time : bool, optional
+        If True, the x-axis is inverted to represent time going backward.
+    default_color : Color, optional
+        The color to use for nodes without the specified metadata.
+    colormap : str | Colormap, optional
+        The colormap to use for coloring continuous values. Defaults to 'viridis'.
+    vmin : float | None, optional
+        The minimum value for normalization. If None, uses the minimum of the data.
+    vmax : float | None, optional
+        The maximum value for normalization. If None, uses the maximum of the data.
+    show_hist : bool, optional
+        Whether to display a histogram of the continuous values.
+    hist_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the histogram.
+    hist_axes_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to define the histogram Axes.
+
+    Returns
+    -------
+    Axes | tuple[Axes, Axes]
+        The Axes with the drawn colored tree,
+        and optionally the histogram Axes if `show_hist` is True.
+    """
+    if show_hist:
+        ax, colors, hist_ax = _init_colored_tree_continuous(
+            tree=tree,
+            color_by=color_by,
+            ax=ax,
+            default_color=default_color,
+            colormap=colormap,
+            vmin=vmin,
+            vmax=vmax,
+            show_hist=show_hist,
+            hist_kwargs=hist_kwargs,
+            hist_axes_kwargs=hist_axes_kwargs,
+        )
+        return draw_tree(
+            tree=tree, ax=ax, colors=colors, backward_time=backward_time
+        ), hist_ax
+
+    ax, colors = _init_colored_tree_continuous(
+        tree=tree,
+        color_by=color_by,
+        ax=ax,
+        default_color=default_color,
+        colormap=colormap,
+        vmin=vmin,
+        vmax=vmax,
+        show_hist=show_hist,
+        hist_kwargs=hist_kwargs,
+        hist_axes_kwargs=hist_axes_kwargs,
+    )
+    return draw_tree(tree=tree, ax=ax, colors=colors, backward_time=backward_time)
+
+
+@overload
+def draw_colored_dated_tree_continuous(
+    tree: Tree,
+    calibration_nodes: tuple[CalibrationNode, CalibrationNode],
+    color_by: str,
+    ax: Axes | None = None,
+    default_color: Color = "black",
+    colormap: str | Colormap = "viridis",
+    vmin: float | None = None,
+    vmax: float | None = None,
+    *,
+    show_hist: Literal[False],
+    hist_kwargs: dict[str, Any] | None = None,
+    hist_axes_kwargs: dict[str, Any] | None = None,
+) -> Axes: ...
+@overload
+def draw_colored_dated_tree_continuous(
+    tree: Tree,
+    calibration_nodes: tuple[CalibrationNode, CalibrationNode],
+    color_by: str,
+    ax: Axes | None = None,
+    default_color: Color = "black",
+    colormap: str | Colormap = "viridis",
+    vmin: float | None = None,
+    vmax: float | None = None,
+    show_hist: Literal[True] = True,
+    hist_kwargs: dict[str, Any] | None = None,
+    hist_axes_kwargs: dict[str, Any] | None = None,
+) -> tuple[Axes, Axes]: ...
+def draw_colored_dated_tree_continuous(
+    tree: Tree,
+    calibration_nodes: tuple[CalibrationNode, CalibrationNode],
+    color_by: str,
+    ax: Axes | None = None,
+    default_color: Color = "black",
+    colormap: str | Colormap = "viridis",
+    vmin: float | None = None,
+    vmax: float | None = None,
+    show_hist: bool = True,
+    hist_kwargs: dict[str, Any] | None = None,
+    hist_axes_kwargs: dict[str, Any] | None = None,
+) -> Axes | tuple[Axes, Axes]:
+    """
+    Draw a dated phylogenetic tree with branches colored based on continuous metadata.
+
+    Parameters
+    ----------
+    tree : Tree
+        The phylogenetic tree to draw.
+    calibration_nodes : tuple[CalibrationNode, CalibrationNode]
+        Two calibration nodes defining the mapping from depth to date.
+    color_by : str
+        The metadata key to color branches by.
+    ax : Axes | None, optional
+        The matplotlib Axes to draw on. If None, uses the current Axes.
+    default_color : Color, optional
+        The color to use for nodes without the specified metadata.
+    colormap : str | Colormap, optional
+        The colormap to use for coloring continuous values. Defaults to 'viridis'.
+    vmin : float | None, optional
+        The minimum value for normalization. If None, uses the minimum of the data.
+    vmax : float | None, optional
+        The maximum value for normalization. If None, uses the maximum of the data.
+    show_hist : bool, optional
+        Whether to display a histogram of the continuous values.
+    hist_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to pass to the histogram.
+    hist_axes_kwargs : dict[str, Any] | None, optional
+        Additional keyword arguments to define the histogram Axes.
+
+    Returns
+    -------
+    Axes | tuple[Axes, Axes]
+        The Axes with the drawn colored dated tree,
+        and optionally the histogram Axes if `show_hist` is True.
+    """
+    if show_hist:
+        ax, colors, hist_ax = _init_colored_tree_continuous(
+            tree=tree,
+            color_by=color_by,
+            ax=ax,
+            default_color=default_color,
+            colormap=colormap,
+            vmin=vmin,
+            vmax=vmax,
+            show_hist=show_hist,
+            hist_kwargs=hist_kwargs,
+            hist_axes_kwargs=hist_axes_kwargs,
+        )
+        return draw_dated_tree(
+            tree=tree,
+            calibration_nodes=calibration_nodes,
+            ax=ax,
+            colors=colors,
+        ), hist_ax
+
+    ax, colors = _init_colored_tree_continuous(
+        tree=tree,
+        color_by=color_by,
+        ax=ax,
+        default_color=default_color,
+        colormap=colormap,
+        vmin=vmin,
+        vmax=vmax,
+        show_hist=show_hist,
+        hist_kwargs=hist_kwargs,
+        hist_axes_kwargs=hist_axes_kwargs,
+    )
+    return draw_dated_tree(
+        tree=tree,
+        calibration_nodes=calibration_nodes,
+        ax=ax,
+        colors=colors,
+    )

phylogenie-3.1.9/src/phylogenie/io/__init__.py

@@ -0,0 +1,3 @@
+from phylogenie.io.fasta import dump_fasta, load_fasta
+
+__all__ = ["load_fasta", "dump_fasta"]

{phylogenie-3.1.5 → phylogenie-3.1.9}/src/phylogenie/io/fasta.py

@@ -1,3 +1,4 @@
+from datetime import date
 from pathlib import Path
 from typing import Callable
 
@@ -5,7 +6,8 @@ from phylogenie.msa import MSA, Sequence
 
 
 def load_fasta(
-    fasta_file: str | Path, extract_time_from_id: Callable[[str], float] | None = None
+    fasta_file: str | Path,
+    extract_time_from_id: Callable[[str], float | date] | None = None,
 ) -> MSA:
     sequences: list[Sequence] = []
     with open(fasta_file, "r") as f:
@@ -17,10 +19,22 @@ def load_fasta(
             if extract_time_from_id is not None:
                 time = extract_time_from_id(id)
             elif "|" in id:
+                last_metadata = id.split("|")[-1]
                 try:
-                    time = float(id.split("|")[-1])
+                    time = float(last_metadata)
                 except ValueError:
-                    pass
+                    try:
+                        time = date.fromisoformat(last_metadata)
+                    except ValueError:
+                        pass
             chars = next(f).strip()
             sequences.append(Sequence(id, chars, time))
     return MSA(sequences)
+
+
+def dump_fasta(msa: MSA | list[Sequence], fasta_file: str | Path) -> None:
+    with open(fasta_file, "w") as f:
+        sequences = msa.sequences if isinstance(msa, MSA) else msa
+        for seq in sequences:
+            f.write(f">{seq.id}\n")
+            f.write(f"{seq.chars}\n")

{phylogenie-3.1.5 → phylogenie-3.1.9}/src/phylogenie/mixins.py

@@ -1,39 +1,31 @@
+from collections.abc import Mapping
 from types import MappingProxyType
-from typing import Any, Mapping, Optional
+from typing import Any
 
 
 class MetadataMixin:
-    """A mixin that provides metadata management with dictionary-like access."""
-
     def __init__(self) -> None:
         self._metadata: dict[str, Any] = {}
 
     @property
     def metadata(self) -> Mapping[str, Any]:
-        """Return a read-only view of all metadata."""
         return MappingProxyType(self._metadata)
 
     def set(self, key: str, value: Any) -> None:
-        """Set or update a metadata value."""
         self._metadata[key] = value
 
     def update(self, metadata: Mapping[str, Any]) -> None:
-        """Bulk update metadata values."""
         self._metadata.update(metadata)
 
-    def get(self, key: str, default: Optional[Any] = None) -> Any:
-        """Get a metadata value, returning `default` if not found."""
+    def get(self, key: str, default: Any = None) -> Any:
         return self._metadata.get(key, default)
 
     def delete(self, key: str) -> None:
-        """Delete a metadata if it exists, else do nothing."""
         self._metadata.pop(key, None)
 
     def clear(self) -> None:
-        """Remove all metadata."""
         self._metadata.clear()
 
-    # Dict-like behavior
     def __getitem__(self, key: str) -> Any:
         return self._metadata[key]
 

{phylogenie-3.1.5 → phylogenie-3.1.9}/src/phylogenie/msa.py

@@ -1,5 +1,6 @@
 from collections.abc import Iterator
 from dataclasses import dataclass
+from datetime import date
 
 import numpy as np
 
@@ -8,7 +9,7 @@ import numpy as np
 class Sequence:
     id: str
     chars: str
-    time: float | None = None
+    time: float | date | None = None
 
 
 class MSA:
@@ -25,8 +26,8 @@ class MSA:
         return [sequence.id for sequence in self.sequences]
 
     @property
-    def times(self) -> list[float]:
-        times: list[float] = []
+    def times(self) -> list[float | date]:
+        times: list[float | date] = []
         for sequence in self:
             if sequence.time is None:
                 raise ValueError(f"Time is not set for sequence {sequence.id}.")

{phylogenie-3.1.5 → phylogenie-3.1.9}/src/phylogenie/treesimulator/__init__.py

@@ -27,11 +27,14 @@ from phylogenie.treesimulator.utils import (
     compute_sackin_index,
     get_distance,
     get_mrca,
+    get_node_ages,
     get_node_depth_levels,
     get_node_depths,
     get_node_height_levels,
     get_node_heights,
     get_node_leaf_counts,
+    get_node_times,
+    get_path,
 )
 
 __all__ = [
@@ -65,9 +68,12 @@ __all__ = [
     "compute_sackin_index",
     "get_distance",
     "get_mrca",
+    "get_node_ages",
     "get_node_depth_levels",
     "get_node_depths",
     "get_node_height_levels",
     "get_node_heights",
     "get_node_leaf_counts",
+    "get_node_times",
+    "get_path",
 ]

{phylogenie-3.1.5 → phylogenie-3.1.9}/src/phylogenie/treesimulator/features.py

@@ -5,11 +5,13 @@ from phylogenie.treesimulator.events.mutations import get_mutation_id
 from phylogenie.treesimulator.model import get_node_state
 from phylogenie.treesimulator.tree import Tree
 from phylogenie.treesimulator.utils import (
+    get_node_ages,
     get_node_depth_levels,
     get_node_depths,
     get_node_height_levels,
     get_node_heights,
     get_node_leaf_counts,
+    get_node_times,
 )
 
 
@@ -22,6 +24,7 @@ def _get_mutations(tree: Tree) -> dict[Tree, int]:
 
 
 class Feature(str, Enum):
+    AGE = "age"
     DEPTH = "depth"
     DEPTH_LEVEL = "depth_level"
     HEIGHT = "height"
@@ -29,9 +32,11 @@ class Feature(str, Enum):
     MUTATION = "mutation"
     N_LEAVES = "n_leaves"
     STATE = "state"
+    TIME = "time"
 
 
 FEATURES_EXTRACTORS = {
+    Feature.AGE: get_node_ages,
     Feature.DEPTH: get_node_depths,
     Feature.DEPTH_LEVEL: get_node_depth_levels,
     Feature.HEIGHT: get_node_heights,
@@ -39,6 +44,7 @@ FEATURES_EXTRACTORS = {
     Feature.MUTATION: _get_mutations,
     Feature.N_LEAVES: get_node_leaf_counts,
     Feature.STATE: _get_states,
+    Feature.TIME: get_node_times,
 }
 
 

{phylogenie-3.1.5 → phylogenie-3.1.9}/src/phylogenie/treesimulator/io/nexus.py

@@ -15,7 +15,7 @@ def _parse_translate_block(lines: Iterator[str]) -> dict[str, str]:
             if ";" in line:
                 return translations
             else:
-                raise ValueError(f"Invalid translate line. Expected '<num> <name>'.")
+                raise ValueError("Invalid translate line. Expected '<num> <name>'.")
         translations[match.group(1)] = match.group(2)
     raise ValueError("Translate block not terminated with ';'.")
 
@@ -33,7 +33,7 @@ def _parse_trees_block(lines: Iterator[str]) -> dict[str, Tree]:
             match = re.match(r"^TREE\s*\*?\s+(\S+)\s*=\s*(.+)$", line, re.IGNORECASE)
             if match is None:
                 raise ValueError(
-                    f"Invalid tree line. Expected 'TREE <name> = <newick>'."
+                    "Invalid tree line. Expected 'TREE <name> = <newick>'."
                 )
             name = match.group(1)
             if name in trees:

{phylogenie-3.1.5 → phylogenie-3.1.9}/src/phylogenie/treesimulator/tree.py

@@ -153,14 +153,27 @@ class Tree(MetadataMixin):
             child.branch_length_or_raise() + child.height for child in self.children
         )
 
+    @property
+    def time(self) -> float:
+        return self.depth
+
+    @property
+    def age(self) -> float:
+        if self.parent is None:
+            return self.height
+        return self.parent.age - self.branch_length_or_raise()
+
     # -------------
     # Miscellaneous
     # -------------
     # Other useful miscellaneous methods.
 
     def ladderize(self, key: Callable[["Tree"], Any] | None = None) -> None:
+        def _default_key(node: Tree) -> int:
+            return node.n_leaves
+
         if key is None:
-            key = lambda node: node.n_leaves
+            key = _default_key
         self._children.sort(key=key)
         for child in self.children:
             child.ladderize(key)
@@ -169,7 +182,7 @@ class Tree(MetadataMixin):
         for node in self:
             if node.name == name:
                 return node
-        raise ValueError(f"Node with name {name} not found.")
+        raise ValueError(f"Node {name} not found.")
 
     def copy(self):
         new_tree = Tree(self.name, self.branch_length)

{phylogenie-3.1.5 → phylogenie-3.1.9}/src/phylogenie/treesimulator/utils.py

@@ -51,6 +51,17 @@ def get_node_heights(tree: Tree) -> dict[Tree, float]:
     return heights
 
 
+def get_node_times(tree: Tree) -> dict[Tree, float]:
+    return get_node_depths(tree)
+
+
+def get_node_ages(tree: Tree) -> dict[Tree, float]:
+    ages: dict[Tree, float] = {tree: tree.height}
+    for node in tree.iter_descendants():
+        ages[node] = ages[node.parent] - node.branch_length  # pyright: ignore
+    return ages
+
+
 def get_mrca(node1: Tree, node2: Tree) -> Tree:
     node1_ancestors = set(node1.iter_upward())
     for node2_ancestor in node2.iter_upward():

{phylogenie-3.1.5 → phylogenie-3.1.9/src/phylogenie.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: phylogenie
-Version: 3.1.5
+Version: 3.1.9
 Summary: Generate phylogenetic datasets with minimal setup effort
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

phylogenie-3.1.5/src/phylogenie/draw.py

@@ -1,152 +0,0 @@
-from enum import Enum
-from typing import Any, Callable
-
-import matplotlib.colors as mcolors
-import matplotlib.patches as mpatches
-import matplotlib.pyplot as plt
-from matplotlib.axes import Axes
-from matplotlib.colors import Colormap
-from mpl_toolkits.axes_grid1.inset_locator import inset_axes  # pyright: ignore
-
-from phylogenie.treesimulator import Tree, get_node_depth_levels, get_node_depths
-
-
-class Coloring(str, Enum):
-    DISCRETE = "discrete"
-    CONTINUOUS = "continuous"
-
-
-Color = str | tuple[float, float, float] | tuple[float, float, float, float]
-
-
-def draw_colored_tree(
-    tree: Tree, ax: Axes | None = None, colors: Color | dict[Tree, Color] = "black"
-) -> Axes:
-    if ax is None:
-        ax = plt.gca()
-
-    if not isinstance(colors, dict):
-        colors = {node: colors for node in tree}
-
-    xs = (
-        get_node_depth_levels(tree)
-        if any(node.branch_length is None for node in tree.iter_descendants())
-        else get_node_depths(tree)
-    )
-    ys: dict[Tree, float] = {node: i for i, node in enumerate(tree.get_leaves())}
-    for node in tree.postorder_traversal():
-        if node.is_internal():
-            ys[node] = sum(ys[child] for child in node.children) / len(node.children)
-
-    if tree.branch_length is not None:
-        ax.hlines(y=ys[tree], xmin=0, xmax=xs[tree], color=colors[tree])  # pyright: ignore
-    for node in tree:
-        x1, y1 = xs[node], ys[node]
-        for child in node.children:
-            x2, y2 = xs[child], ys[child]
-            ax.hlines(y=y2, xmin=x1, xmax=x2, color=colors[child])  # pyright: ignore
-            ax.vlines(x=x1, ymin=y1, ymax=y2, color=colors[child])  # pyright: ignore
-
-    ax.set_yticks([])  # pyright: ignore
-    return ax
-
-
-def draw_tree(
-    tree: Tree,
-    ax: Axes | None = None,
-    color_by: str | dict[str, Any] | None = None,
-    coloring: str | Coloring | None = None,
-    default_color: Color = "black",
-    colormap: str | Colormap | None = None,
-    vmin: float | None = None,
-    vmax: float | None = None,
-    show_legend: bool = True,
-    labels: dict[Any, Any] | None = None,
-    legend_kwargs: dict[str, Any] | None = None,
-    show_hist: bool = True,
-    hist_kwargs: dict[str, Any] | None = None,
-    hist_axes_kwargs: dict[str, Any] | None = None,
-) -> Axes | tuple[Axes, Axes]:
-    if ax is None:
-        ax = plt.gca()
-
-    if color_by is None:
-        return draw_colored_tree(tree, ax, colors=default_color)
-
-    if isinstance(color_by, str):
-        features = {node: node[color_by] for node in tree if color_by in node.metadata}
-    else:
-        features = {node: color_by[node.name] for node in tree if node.name in color_by}
-    values = list(features.values())
-
-    if coloring is None:
-        coloring = (
-            Coloring.CONTINUOUS
-            if any(isinstance(f, float) for f in values)
-            else Coloring.DISCRETE
-        )
-    if colormap is None:
-        colormap = "tab20" if coloring == Coloring.DISCRETE else "viridis"
-    if isinstance(colormap, str):
-        colormap = plt.get_cmap(colormap)
-
-    def _get_colors(feature_map: Callable[[Any], Color]) -> dict[Tree, Color]:
-        return {
-            node: feature_map(features[node]) if node in features else default_color
-            for node in tree
-        }
-
-    if coloring == Coloring.DISCRETE:
-        if any(isinstance(f, float) for f in values):
-            raise ValueError(
-                "Discrete coloring selected but feature values are not all categorical."
-            )
-        feature_colors = {
-            f: mcolors.to_hex(colormap(i)) for i, f in enumerate(set(values))
-        }
-        colors = _get_colors(lambda f: feature_colors[f])
-
-        if show_legend:
-            legend_handles = [
-                mpatches.Patch(
-                    color=feature_colors[f],
-                    label=str(f) if labels is None else labels[f],
-                )
-                for f in feature_colors
-            ]
-            if any(color_by not in node.metadata for node in tree):
-                legend_handles.append(mpatches.Patch(color=default_color, label="NA"))
-            if legend_kwargs is None:
-                legend_kwargs = {}
-            ax.legend(handles=legend_handles, **legend_kwargs)  # pyright: ignore
-
-        return draw_colored_tree(tree, ax, colors)
-
-    if coloring == Coloring.CONTINUOUS:
-        vmin = min(values) if vmin is None else vmin
-        vmax = max(values) if vmax is None else vmax
-        norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
-        colors = _get_colors(lambda f: colormap(norm(float(f))))
-
-        if show_hist:
-            default_hist_axes_kwargs = {"width": "25%", "height": "25%"}
-            if hist_axes_kwargs is not None:
-                default_hist_axes_kwargs.update(hist_axes_kwargs)
-            hist_ax = inset_axes(ax, **default_hist_axes_kwargs)  # pyright: ignore
-
-            hist_kwargs = {} if hist_kwargs is None else hist_kwargs
-            _, bins, patches = hist_ax.hist(values, **hist_kwargs)  # pyright: ignore
-
-            for patch, b0, b1 in zip(patches, bins[:-1], bins[1:]):  # pyright: ignore
-                midpoint = (b0 + b1) / 2  # pyright: ignore
-                patch.set_facecolor(colormap(norm(midpoint)))  # pyright: ignore
-            return draw_colored_tree(tree, ax, colors), hist_ax  # pyright: ignore
-
-        else:
-            sm = plt.cm.ScalarMappable(cmap=colormap, norm=norm)
-            ax.get_figure().colorbar(sm, ax=ax)  # pyright: ignore
-            return draw_colored_tree(tree, ax, colors)
-
-    raise ValueError(
-        f"Unknown coloring method: {coloring}. Choices are {list(Coloring)}."
-    )

phylogenie-3.1.5/src/phylogenie/io/__init__.py

@@ -1,3 +0,0 @@
-from phylogenie.io.fasta import load_fasta
-
-__all__ = ["load_fasta"]