phylogenie 3.1.11__py3-none-any.whl → 3.1.13__py3-none-any.whl

phylogenie/draw.py

@@ -224,7 +224,7 @@ def _init_colored_tree_categorical(
     color_by: str,
     ax: Axes | None = None,
     default_color: Color = "black",
-    colormap: str | Colormap = "tab20",
+    colormap: str | dict[str, Color] | Colormap = "tab20",
     show_legend: bool = True,
     labels: dict[Any, str] | None = None,
     legend_kwargs: dict[str, Any] | None = None,
@@ -242,8 +242,11 @@ def _init_colored_tree_categorical(
         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'.
+    colormap : str | dict[str, Color] | Colormap, optional
+        The colormap to use for coloring categories.
+        If a string, it is used to get a matplotlib colormap.
+        If a dict, it maps category values to colors directly.
+        Defaults to 'tab20'.
     show_legend : bool, optional
         Whether to display a legend for the categories.
     labels : dict[Any, str] | None, optional
@@ -259,13 +262,16 @@ def _init_colored_tree_categorical(
     if ax is None:
         ax = plt.gca()
 
+    features = {node: node[color_by] for node in tree if color_by in node.metadata}
     if isinstance(colormap, str):
         colormap = plt.get_cmap(colormap)
+    if isinstance(colormap, Colormap):
+        feature_colors = {
+            f: mcolors.to_hex(colormap(i)) for i, f in enumerate(set(features.values()))
+        }
+    else:
+        feature_colors = 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
@@ -294,7 +300,7 @@ def draw_colored_tree_categorical(
     ax: Axes | None = None,
     backward_time: bool = False,
     default_color: Color = "black",
-    colormap: str | Colormap = "tab20",
+    colormap: str | dict[str, Color] | Colormap = "tab20",
     show_legend: bool = True,
     labels: dict[Any, str] | None = None,
     legend_kwargs: dict[str, Any] | None = None,
@@ -316,8 +322,11 @@ def draw_colored_tree_categorical(
         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'.
+    colormap : str | dict[str, Color] | Colormap, optional
+        The colormap to use for coloring categories.
+        If a string, it is used to get a matplotlib colormap.
+        If a dict, it maps category values to colors directly.
+        Defaults to 'tab20'.
     show_legend : bool, optional
         Whether to display a legend for the categories.
     labels : dict[Any, str] | None, optional
@@ -360,7 +369,7 @@ def draw_colored_dated_tree_categorical(
     color_by: str,
     ax: Axes | None = None,
     default_color: Color = "black",
-    colormap: str | Colormap = "tab20",
+    colormap: str | dict[str, Color] | Colormap = "tab20",
     show_legend: bool = True,
     labels: dict[Any, str] | None = None,
     legend_kwargs: dict[str, Any] | None = None,
@@ -381,8 +390,11 @@ def draw_colored_dated_tree_categorical(
         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'.
+    colormap : str | dict[str, Color] | Colormap, optional
+        The colormap to use for coloring categories.
+        If a string, it is used to get a matplotlib colormap.
+        If a dict, it maps category values to colors directly.
+        Defaults to 'tab20'.
     show_legend : bool, optional
         Whether to display a legend for the categories.
     labels : dict[Any, str] | None, optional

{phylogenie-3.1.11.dist-info → phylogenie-3.1.13.dist-info}/METADATA

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

{phylogenie-3.1.11.dist-info → phylogenie-3.1.13.dist-info}/RECORD

@@ -1,5 +1,5 @@
 phylogenie/__init__.py,sha256=RyH_mcEXcKtELTkl9HRzWcl2ANWN-G4K0BLigqR3_B4,2888
-phylogenie/draw.py,sha256=-hyycRN4RyG0nzZrE4CC4bmg0PgxJOeK67rwlmFZiuc,26841
+phylogenie/draw.py,sha256=z0f5g8t54DST-AxvFyNYREqpj41IhSIAjHIdrtVRH08,27449
 phylogenie/main.py,sha256=ry3B3HiwibZG3_qB58T5UhWy5dp6neYUtSqzL9LrSkA,1698
 phylogenie/mixins.py,sha256=wMwqP6zkqME9eMyzx5FS6-p9X8yW09jIC8jge8pHlkk,907
 phylogenie/msa.py,sha256=RGXxmo1WX9wi5NpEQTxMuH4zk8vM0f8A7Knm0YUJ9AI,2111
@@ -33,9 +33,9 @@ phylogenie/treesimulator/events/mutations.py,sha256=8Nqa2fg7fwaVNe5XSkGDSwp9pIKQ
 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=zqT9dzj413z_s0hqp3Cdq5NMO6lv-zuuaJlaqzaqaB8,1847
-phylogenie-3.1.11.dist-info/licenses/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-3.1.11.dist-info/METADATA,sha256=91J5lodQcF09WeD0Oq8ljaIVEWQzoYeXvb2VDbNtQ7U,5200
-phylogenie-3.1.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-phylogenie-3.1.11.dist-info/entry_points.txt,sha256=BBH8LoReHnNFnvq4sROEsVFegfkKJ6c_oHZ7bgK7Jl4,52
-phylogenie-3.1.11.dist-info/top_level.txt,sha256=1YGZJhKA9tN9qI0Hcj6Cn_sOoDpba0HQlNcgQTjMD-8,11
-phylogenie-3.1.11.dist-info/RECORD,,
+phylogenie-3.1.13.dist-info/licenses/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-3.1.13.dist-info/METADATA,sha256=k0zK6rE7YTrt5ulq4i8doSKb-TH5GqoGdKmFt-Zam8c,5200
+phylogenie-3.1.13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+phylogenie-3.1.13.dist-info/entry_points.txt,sha256=BBH8LoReHnNFnvq4sROEsVFegfkKJ6c_oHZ7bgK7Jl4,52
+phylogenie-3.1.13.dist-info/top_level.txt,sha256=1YGZJhKA9tN9qI0Hcj6Cn_sOoDpba0HQlNcgQTjMD-8,11
+phylogenie-3.1.13.dist-info/RECORD,,