phylogenie 3.1.4__py3-none-any.whl → 3.1.6__py3-none-any.whl

phylogenie/draw.py

@@ -1,7 +1,9 @@
-from enum import Enum
-from typing import Any, Callable
+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
@@ -11,15 +13,43 @@ 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"
+@dataclass
+class CalibrationNode:
+    node: Tree
+    date: datetime
 
 
 Color = str | tuple[float, float, float] | tuple[float, float, float, float]
 
 
-def _draw_colored_tree(tree: Tree, ax: Axes, colors: Color | dict[Tree, Color]) -> Axes:
+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}
 
@@ -28,122 +58,633 @@ def _draw_colored_tree(tree: Tree, ax: Axes, colors: Color | dict[Tree, Color])
         if any(node.branch_length is None for node in tree.iter_descendants())
         else get_node_depths(tree)
     )
+    if backward_time:
+        max_x = max(xs.values())
+        xs = {node: max_x - x for node, x in xs.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:
+        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]
-        if node.parent is None:
-            ax.hlines(y=y1, xmin=0, xmax=x1, color=colors[node])  # pyright: ignore
-            continue
-        x0, y0 = xs[node.parent], ys[node.parent]
-        ax.vlines(x=x0, ymin=y0, ymax=y1, color=colors[node])  # pyright: ignore
-        ax.hlines(y=y1, xmin=x0, xmax=x1, color=colors[node])  # pyright: ignore
+        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 draw_tree(
+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,
-    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,
+    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, Any] | None = None,
+    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],
+    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,
-) -> Axes | tuple[Axes, Axes]:
+) -> 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 color_by is None:
-        return _draw_colored_tree(tree, ax, colors=default_color)
+    if isinstance(colormap, str):
+        colormap = plt.get_cmap(colormap)
 
-    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}
+    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],
+    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.
 
-    if coloring is None:
-        coloring = (
-            Coloring.CONTINUOUS
-            if any(isinstance(f, float) for f in values)
-            else Coloring.DISCRETE
+    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,
         )
-    if colormap is None:
-        colormap = "tab20" if coloring == Coloring.DISCRETE else "viridis"
-    if isinstance(colormap, str):
-        colormap = plt.get_cmap(colormap)
+        return draw_tree(
+            tree=tree, ax=ax, colors=colors, backward_time=backward_time
+        ), hist_ax
 
-    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
-        }
+    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)
 
-    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(  # pyright: ignore
-                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)}."
+
+@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],
+    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/io/fasta.py

@@ -24,3 +24,11 @@ def load_fasta(
             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/main.py

@@ -2,23 +2,44 @@ import os
 from argparse import ArgumentParser
 from glob import glob
 
-from pydantic import TypeAdapter
+from pydantic import TypeAdapter, ValidationError
 from yaml import safe_load
 
 from phylogenie.generators import DatasetGeneratorConfig
 from phylogenie.generators.dataset import DatasetGenerator
 
 
-def run(config_path: str) -> None:
+def _format_validation_error(e: ValidationError) -> str:
+    formatted_errors = [
+        f"- {'.'.join(str(loc) for loc in err['loc'])}: {err['msg']} ({err['type']})"
+        for err in e.errors()
+    ]
+    return "\n".join(formatted_errors)
+
+
+def _generate_from_config_file(config_file: str):
     adapter: TypeAdapter[DatasetGenerator] = TypeAdapter(DatasetGeneratorConfig)
+    with open(config_file, "r") as f:
+        try:
+            config = safe_load(f)
+        except Exception as e:
+            print(f"❌ Failed to parse {config_file}: {e}")
+            exit(-1)
+        try:
+            generator = adapter.validate_python(config)
+        except ValidationError as e:
+            print("❌ Invalid configuration:")
+            print(_format_validation_error(e))
+            exit(-1)
+    generator.generate()
 
+
+def run(config_path: str) -> None:
     if os.path.isdir(config_path):
         for config_file in glob(os.path.join(config_path, "**/*.yaml"), recursive=True):
-            with open(config_file, "r") as f:
-                adapter.validate_python(safe_load(f)).generate()
+            _generate_from_config_file(config_file)
     else:
-        with open(config_path, "r") as f:
-            adapter.validate_python(safe_load(f)).generate()
+        _generate_from_config_file(config_path)
 
 
 def main() -> None:

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/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/treesimulator/tree.py

@@ -159,8 +159,11 @@ class Tree(MetadataMixin):
     # 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 +172,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.4.dist-info → phylogenie-3.1.6.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: phylogenie
-Version: 3.1.4
-Summary: Add your description here
+Version: 3.1.6
+Summary: Generate phylogenetic datasets with minimal setup effort
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt

{phylogenie-3.1.4.dist-info → phylogenie-3.1.6.dist-info}/RECORD

@@ -1,7 +1,7 @@
 phylogenie/__init__.py,sha256=-4GKO2EKgc2GiOhzVRI0SQk6wQZ4AlP7Ron_Lp-s_r0,2913
-phylogenie/draw.py,sha256=nT3FOHSm85G23qz2xxYf2MesEDonWRSxfKGJkVboTKY,5580
-phylogenie/main.py,sha256=DstKWUySyDXdcPN_P8O3Mmoau4IjUxCQ4iLWDVuAcws,1054
-phylogenie/mixins.py,sha256=kwxcHGYSS-Frk-r0nwuyILYS1g6NjSOhbUKEYt2g1B0,1319
+phylogenie/draw.py,sha256=djG0cG9mmNFZ7vf3jo8Ei3bNBJVqujp6gtZtwPxyfSY,23341
+phylogenie/main.py,sha256=ry3B3HiwibZG3_qB58T5UhWy5dp6neYUtSqzL9LrSkA,1698
+phylogenie/mixins.py,sha256=wMwqP6zkqME9eMyzx5FS6-p9X8yW09jIC8jge8pHlkk,907
 phylogenie/msa.py,sha256=JDGyZUsAq6-m-SQjoCDjAkAZIxfgyl_PDIhdYn5HOow,2064
 phylogenie/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
@@ -14,7 +14,7 @@ phylogenie/generators/factories.py,sha256=2mTFdFbbLyV3v79JaOEVtqLOmxQHaOUv1S-Y3v
 phylogenie/generators/trees.py,sha256=8dO1CkU34E6mmMAHrYqiLV_VA8r54cSEOo-UzoHiN20,10467
 phylogenie/generators/typeguards.py,sha256=yj4VkhOaUXJ2OrY-6zhOeY9C4yKIQxjZtk2d-vIxttQ,828
 phylogenie/io/__init__.py,sha256=eWDU6YDqAdx6TZlUMmMbsfO8gP3i5HP_cfUe_-0x2FA,69
-phylogenie/io/fasta.py,sha256=CUFO06m7wClprarsMheZojM4Os2NQf3ALYUXSWzfNL0,869
+phylogenie/io/fasta.py,sha256=kx9uVATLpzpXdhhNNMvMpB5Vdwh9CTWepTGachvEwV4,1154
 phylogenie/skyline/__init__.py,sha256=7pF4CUb4ZCLzNYJNhOjpuTOLTRhlK7L6ugfccNqjIGo,620
 phylogenie/skyline/matrix.py,sha256=v4SitY7VbXprqlqQckjWTzW5hwRmCyIF595R6IJMxWw,9268
 phylogenie/skyline/parameter.py,sha256=TVqkqirGXNN-VP8hnIJACPkOxUan6LkGa5o_JcPfwbY,4834
@@ -23,7 +23,7 @@ phylogenie/treesimulator/__init__.py,sha256=Qz__cFookXuEjq8AUp1A6XRON0qQ_xX-q2Q5
 phylogenie/treesimulator/features.py,sha256=XbuwGw8xjGs2lNhJvvUUvXVtheSTBaSN6qj39tWYEro,1391
 phylogenie/treesimulator/gillespie.py,sha256=ey2hdpJOSpNW88duwK7wTAdYSTnSuTSZ_yhZv9MlNHo,5323
 phylogenie/treesimulator/model.py,sha256=L0RsL6H1ynFDPecULniSs4Cs8dvz87ovviQOXFy5Qt0,4580
-phylogenie/treesimulator/tree.py,sha256=-yMW14018x9dw45TonS6nRlzWcwXcaHv3Jn5HYriLQQ,6009
+phylogenie/treesimulator/tree.py,sha256=DEdzCh4vABq2f095beh3tD3_aee7EyXPDSjcyHKgKLg,6064
 phylogenie/treesimulator/utils.py,sha256=OxZwVHxN004Jf-kYZ_GfJgIY0beo-0tYq80CuFGQt-M,3416
 phylogenie/treesimulator/events/__init__.py,sha256=w2tJ0D2WB5AiCbr3CsKN6vdADueiAEMzd_ve0rpa4zg,939
 phylogenie/treesimulator/events/base.py,sha256=JQKYUZmhB2Q-WQOy2ULGKQiabsMz-JvwMVfDoa3ZKyo,1170
@@ -32,10 +32,10 @@ phylogenie/treesimulator/events/core.py,sha256=bhgQgi5L-oaHsoWJmUOsTTzWxi0POYxVL
 phylogenie/treesimulator/events/mutations.py,sha256=8Nqa2fg7fwaVNe5XSkGDSwp9pIKQ7XaBQCCj-LYlfzA,3666
 phylogenie/treesimulator/io/__init__.py,sha256=rfP-zp8SP8baq5_4dPAr10WH0W6KfoMCxdTZDCSXtzE,185
 phylogenie/treesimulator/io/newick.py,sha256=8Pr_jixByPOaVch18w-rFt62HYy0U97YMu0H-QSwIy0,3449
-phylogenie/treesimulator/io/nexus.py,sha256=tPZRPejnG0OEE7YTorwvyZe2vnZHv4KiC9f-N7Kx6qg,1849
-phylogenie-3.1.4.dist-info/licenses/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-3.1.4.dist-info/METADATA,sha256=en3CvCPhiUbQPZTAUYWpx8nQze6oqS-KQJJa3xGdvbI,5163
-phylogenie-3.1.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-phylogenie-3.1.4.dist-info/entry_points.txt,sha256=BBH8LoReHnNFnvq4sROEsVFegfkKJ6c_oHZ7bgK7Jl4,52
-phylogenie-3.1.4.dist-info/top_level.txt,sha256=1YGZJhKA9tN9qI0Hcj6Cn_sOoDpba0HQlNcgQTjMD-8,11
-phylogenie-3.1.4.dist-info/RECORD,,
+phylogenie/treesimulator/io/nexus.py,sha256=zqT9dzj413z_s0hqp3Cdq5NMO6lv-zuuaJlaqzaqaB8,1847
+phylogenie-3.1.6.dist-info/licenses/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-3.1.6.dist-info/METADATA,sha256=az3fV29mfXpVsyxK5PDssutWcvaLnu_MLfHO4pgyBas,5194
+phylogenie-3.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+phylogenie-3.1.6.dist-info/entry_points.txt,sha256=BBH8LoReHnNFnvq4sROEsVFegfkKJ6c_oHZ7bgK7Jl4,52
+phylogenie-3.1.6.dist-info/top_level.txt,sha256=1YGZJhKA9tN9qI0Hcj6Cn_sOoDpba0HQlNcgQTjMD-8,11
+phylogenie-3.1.6.dist-info/RECORD,,