phylogenie 2.1.8__py3-none-any.whl → 2.1.9__py3-none-any.whl

phylogenie/plot.py

@@ -1,4 +1,7 @@
+from enum import Enum
+
 import matplotlib.colors as mcolors
+import matplotlib.patches as mpatches
 import matplotlib.pyplot as plt
 
 from phylogenie import Tree
@@ -6,34 +9,69 @@ from phylogenie.tree import Tree
 from phylogenie.utils import get_times
 
 
+class Coloring(str, Enum):
+    DISCRETE = "discrete"
+    CONTINUOUS = "continuous"
+
+
 def plot_tree(
     tree: Tree,
     ax: plt.Axes | None = None,  # pyright: ignore
     color_by: str | None = None,
     default_color: str = "black",
-    cmap: str = "tab20",
+    coloring: str | Coloring | None = None,
+    cmap: str | None = None,
+    show_legend: bool = True,
 ) -> plt.Axes:  # pyright: ignore
     if ax is None:
         ax = plt.gca()
 
     xs = get_times(tree)
-    ys = {node.name: i for i, node in enumerate(tree.inorder_traversal())}
+    ys = {node: i for i, node in enumerate(tree.inorder_traversal())}
+
     if color_by is not None:
         features = set(node.get(color_by) for node in tree)
-        feature_colors = {
-            f: mcolors.to_hex(plt.get_cmap(cmap, len(features))(i))
-            for i, f in enumerate(features)
-        }
-        colors = {node.name: feature_colors[node.get(color_by)] for node in tree}
+        if coloring is None and any(isinstance(f, float) for f in features):
+            coloring = Coloring.CONTINUOUS
+        elif coloring is None:
+            coloring = Coloring.DISCRETE
+        if coloring == Coloring.DISCRETE:
+            if any(isinstance(f, float) for f in features):
+                raise ValueError(
+                    "Discrete coloring selected but feature values are not all categorical."
+                )
+            cmap = "tab20" if cmap is None else cmap
+            colormap = plt.get_cmap(cmap, len(features))
+            feature_colors = {
+                f: mcolors.to_hex(colormap(i)) for i, f in enumerate(features)
+            }
+            colors = {node: feature_colors[node.get(color_by)] for node in tree}
+            legend_handles = [
+                mpatches.Patch(color=feature_colors[f], label=str(f)) for f in features
+            ]
+            if show_legend:
+                ax.legend(handles=legend_handles, title=color_by)  # pyright: ignore
+        elif coloring in {Coloring.CONTINUOUS}:
+            cmap = "viridis" if cmap is None else cmap
+            values = list(map(float, features))
+            norm = mcolors.Normalize(vmin=min(values), vmax=max(values))
+            colormap = plt.get_cmap(cmap)
+            colors = {node: colormap(norm(float(node.get(color_by)))) for node in tree}
+        else:
+            raise ValueError(
+                f"Unknown coloring method: {coloring}. Choices are {list(Coloring)}."
+            )
     else:
-        colors = {node.name: default_color for node in tree}
+        colors = {node: default_color for node in tree}
 
     for node in tree:
+        x1, y1 = xs[node.name], 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.name], ys[node.parent.name]
-        x1, y1 = xs[node.name], ys[node.name]
-        ax.plot([x0, x0], [y0, y1], color=colors[node.name])  # pyright: ignore
-        ax.plot([x0, x1], [y1, y1], color=colors[node.name])  # pyright: ignore
+        x0, y0 = xs[node.parent.name], 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
+
     ax.set_yticks([])  # pyright: ignore
     return ax

{phylogenie-2.1.8.dist-info → phylogenie-2.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.1.8
+Version: 2.1.9
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com

{phylogenie-2.1.8.dist-info → phylogenie-2.1.9.dist-info}/RECORD

@@ -10,7 +10,7 @@ phylogenie/io.py,sha256=nwy8DOknt0HqF9qMeFZHrCmSXpM5AGrU5oajwTtD6vY,3973
 phylogenie/main.py,sha256=vtvSpQxBNlYABoFQ25czl-l3fIr4QRo3svWVd-jcArw,1170
 phylogenie/models.py,sha256=pCg9ob0RpLUHwM49x4knKxL4FNPr3-EU_6zMXsvxtAg,370
 phylogenie/msa.py,sha256=JDGyZUsAq6-m-SQjoCDjAkAZIxfgyl_PDIhdYn5HOow,2064
-phylogenie/plot.py,sha256=yLJ331UNta0J3d8qdLBtLdIXsC3Qd0anA1qf_5AUwK8,1310
+phylogenie/plot.py,sha256=9lXZokOi7sJ1HLDPQArzjODnQLL_nv16Fr7AtP-5eFE,2869
 phylogenie/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phylogenie/skyline/__init__.py,sha256=7pF4CUb4ZCLzNYJNhOjpuTOLTRhlK7L6ugfccNqjIGo,620
 phylogenie/skyline/matrix.py,sha256=Gl8OgKjtieG0NwPYiPimKI36gefV8fm_OeorjdXxPTs,9146
@@ -28,8 +28,8 @@ phylogenie/treesimulator/model.py,sha256=Ct0lfn6maKtjuFxivWx1MbFHvH3Y-fiJ0XXMdkN
 phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
 phylogenie/typings.py,sha256=GknvAFXyiaWeeYJ8Lk5d6E2VHT-xW6ONEojYbtJYiB8,476
 phylogenie/utils.py,sha256=JDe6_8-9AFYf7SuTrqUXTZ7n3CA-goQGg9vTi5UcHB8,878
-phylogenie-2.1.8.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-2.1.8.dist-info/METADATA,sha256=-F1xhii9rWBPtXGyjIr0UTVWbsrD7wnswCbw--LoS8I,5476
-phylogenie-2.1.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-phylogenie-2.1.8.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
-phylogenie-2.1.8.dist-info/RECORD,,
+phylogenie-2.1.9.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.1.9.dist-info/METADATA,sha256=V-aDacuPL0dfohsDmU69cvzHakcf6RfdDvrgK3zSCgc,5476
+phylogenie-2.1.9.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.1.9.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.1.9.dist-info/RECORD,,