iplotx 0.0.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

iplotx/plotting.py

@@ -1,4 +1,6 @@
-from typing import Union, Sequence
+from typing import Optional, Sequence
+from contextlib import nullcontext
+import numpy as np
 import pandas as pd
 import matplotlib as mpl
 import matplotlib.pyplot as plt
@@ -7,77 +9,176 @@ from .typing import (
     GraphType,
     LayoutType,
     GroupingType,
+    TreeType,
 )
 from .network import NetworkArtist
 from .groups import GroupingArtist
-from .styles import stylecontext
-
-
-def plot(
-    network: Union[GraphType, None] = None,
-    layout: Union[LayoutType, None] = None,
-    grouping: Union[None, GroupingType] = None,
-    vertex_labels: Union[None, list, dict, pd.Series] = None,
-    edge_labels: Union[None, Sequence] = None,
-    ax: Union[None, object] = None,
-    styles: Sequence[Union[str, dict]] = (),
-):
-    """Plot this network using the specified layout.
+from .tree import TreeArtist
+from .style import context
+
+
+def network(
+    network: Optional[GraphType] = None,
+    layout: Optional[LayoutType] = None,
+    grouping: Optional[GroupingType] = None,
+    vertex_labels: Optional[list | dict | pd.Series] = None,
+    edge_labels: Optional[Sequence] = None,
+    ax: Optional[mpl.axes.Axes] = None,
+    style: str | dict | Sequence[str | dict] = (),
+    title: Optional[str] = None,
+    aspect: Optional[str | float] = None,
+    margins: float | tuple[float, float] = 0,
+    **kwargs,
+) -> list[mpl.artist.Artist]:
+    """Plot this network and/or vertex grouping using the specified layout.
 
     Parameters:
-        network (GraphType): The network to plot. Can be a networkx or igraph graph.
-        layout (Union[LayoutType, None], optional): The layout to use for plotting. If None, a layout will be looked for in the network object and, if none is found, an exception is raised. Defaults to None.
-        vertex_labels (list, dict, or pandas.Series): The labels for the vertices. If None, no vertex labels
+        network: The network to plot. Can be a networkx or igraph graph.
+        layout: The layout to use for plotting. If None, a layout will be looked for in the
+            network object and, if none is found, an exception is raised. Defaults to None.
+        vertex_labels: The labels for the vertices. If None or False, no vertex labels
             will be drawn. If a list, the labels are taken from the list. If a dict, the keys
-            should be the vertex IDs and the values should be the labels.
-        edge_labels (Union[None, Sequence], optional): The labels for the edges. If None, no edge labels will be drawn. Defaults to None.
-        ax (Union[None, object], optional): The axis to plot on. If None, a new figure and axis will be created. Defaults to None.
-        **style: Additional keyword arguments are treated as style for the objects to plot.
+            should be the vertex IDs and the values should be the labels. If True (a single
+            bool value), the vertex IDs will be used as labels.
+        edge_labels: The labels for the edges. If None, no edge labels will be drawn. Defaults
+            to None.
+        ax: The axis to plot on. If None, a new figure and axis will be created. Defaults to
+            None.
+        style: Apply this style for the objects to plot. This can be a sequence (e.g. list)
+            of styles and they will be applied in order.
+        title: If not None, set the axes title to this value.
+        aspect: If not None, set the aspect ratio of the axis to this value. The most common
+            value is 1.0, which proportionates x- and y-axes.
+        margins: How much margin to leave around the plot. A higher value (e.g. 0.1) can be
+            used as a quick fix when some vertex shapes reach beyond the plot edge. This is
+            a fraction of the data limits, so 0.1 means 10% of the data limits will be left
+            as margin.
+        **kwargs: Additional arguments are treated as an alternate way to specify style. If
+            both "style" and additional **kwargs are provided, they are both applied in that
+            order (style, then **kwargs).
 
     Returns:
-        A NetworkArtist object.
+        A list of mpl.artist.Artist objects, set as a direct child of the matplotlib Axes.
+        The list can have one or two elements, depending on whether you are requesting to
+        plot a network, a grouping, or both.
     """
-    if len(styles):
-        with stylecontext(styles):
-            return plot(
-                network=network,
-                layout=layout,
-                grouping=grouping,
+    stylecontext = context(style, **kwargs) if style or kwargs else nullcontext()
+
+    with stylecontext:
+        if (network is None) and (grouping is None):
+            raise ValueError("At least one of network or grouping must be provided.")
+
+        if ax is None:
+            fig, ax = plt.subplots()
+
+        artists = []
+        if network is not None:
+            nwkart = NetworkArtist(
+                network,
+                layout,
+                vertex_labels=vertex_labels,
                 edge_labels=edge_labels,
+                transform=mpl.transforms.IdentityTransform(),
+                offset_transform=ax.transData,
+            )
+            ax.add_artist(nwkart)
+
+            # Set the figure, which itself sets the dpi scale for vertices, edges,
+            # arrows, etc. Now data limits can be computed correctly
+            nwkart.set_figure(ax.figure)
+
+            artists.append(nwkart)
+
+            # Set normailsed layout since we have it by now
+            layout = nwkart.get_layout()
+
+        if grouping is not None:
+            grpart = GroupingArtist(
+                grouping,
+                layout,
+                network=network,
+                transform=ax.transData,
             )
+            ax.add_artist(grpart)
+
+            grpart.set_figure(ax.figure)
+            artists.append(grpart)
+
+        if title is not None:
+            ax.set_title(title)
 
-    if (network is None) and (grouping is None):
-        raise ValueError("At least one of network or grouping must be provided.")
+        if aspect is not None:
+            ax.set_aspect(aspect)
 
-    if ax is None:
-        fig, ax = plt.subplots()
+        _postprocess_axis(ax, artists)
 
-    artists = []
-    if network is not None:
-        nwkart = NetworkArtist(
-            network,
-            layout,
+        if np.isscalar(margins):
+            margins = (margins, margins)
+        if (margins[0] != 0) or (margins[1] != 0):
+            ax.margins(*margins)
+
+        return artists
+
+
+def tree(
+    tree: Optional[TreeType] = None,
+    layout: str | LayoutType = "horizontal",
+    orientation: str = "right",
+    directed: bool | str = False,
+    vertex_labels: Optional[list | dict | pd.Series] = None,
+    ax: Optional[mpl.axes.Axes] = None,
+    style: str | dict | Sequence[str | dict] = "tree",
+    title: Optional[str] = None,
+    aspect: Optional[str | float] = None,
+    margins: float | tuple[float, float] = 0,
+    **kwargs,
+) -> TreeArtist:
+    """Plot a tree using the specified layout.
+
+    Parameters:
+        tree: The tree to plot. Can be a BioPython.Phylo.Tree object.
+        layout: The layout to use for plotting.
+        orientation: The orientation of the horizontal layout. Can be "right" or "left". Defaults to
+            "right".
+        directed: If False, donot draw arrows. If True or "child", draw arrows from parent to child
+            node. If "parent", draw arrows the other way around.
+
+    Returns:
+        A TreeArtist object, set as a direct child of the matplotlib Axes.
+    """
+    stylecontext = context(style, **kwargs) if style or kwargs else nullcontext()
+
+    with stylecontext:
+        if ax is None:
+            fig, ax = plt.subplots()
+
+        artist = TreeArtist(
+            tree=tree,
+            layout=layout,
+            orientation=orientation,
+            directed=directed,
+            transform=mpl.transforms.IdentityTransform(),
+            offset_transform=ax.transData,
             vertex_labels=vertex_labels,
-            edge_labels=edge_labels,
-        )
-        ax.add_artist(nwkart)
-        # Postprocess for things that require an axis (transform, etc.)
-        nwkart._process()
-        artists.append(nwkart)
-
-    if grouping is not None:
-        grpart = GroupingArtist(
-            grouping,
-            layout,
         )
-        ax.add_artist(grpart)
-        # Postprocess for things that require an axis (transform, etc.)
-        grpart._process()
-        artists.append(grpart)
+        ax.add_artist(artist)
+
+        artist.set_figure(ax.figure)
+
+        if title is not None:
+            ax.set_title(title)
+
+        if aspect is not None:
+            ax.set_aspect(aspect)
+
+        _postprocess_axis(ax, [artist])
 
-    _postprocess_axis(ax, artists)
+        if np.isscalar(margins):
+            margins = (margins, margins)
+        if (margins[0] != 0) or (margins[1] != 0):
+            ax.margins(*margins)
 
-    return artists
+    return artist
 
 
 # INTERNAL ROUTINES
@@ -98,7 +199,8 @@ def _postprocess_axis(ax, artists):
     bboxes = []
     for art in artists:
         bboxes.append(art.get_datalim(ax.transData))
-    ax.update_datalim(mpl.transforms.Bbox.union(bboxes))
+    bbox = mpl.transforms.Bbox.union(bboxes)
+    ax.update_datalim(bbox)
 
     # Autoscale for x/y axis limits
     ax.autoscale_view()

iplotx/style.py

@@ -0,0 +1,379 @@
+from typing import (
+    Any,
+    Never,
+    Optional,
+    Sequence,
+)
+from collections.abc import Hashable
+import copy
+from contextlib import contextmanager
+import numpy as np
+import pandas as pd
+
+
+style_leaves = (
+    "cmap",
+    "color",
+    "size",
+    "edgecolor",
+    "facecolor",
+    "linewidth",
+    "linestyle",
+    "alpha",
+    "zorder",
+    "tension",
+    "looptension",
+    "lloopmaxangle",
+    "rotate",
+    "marker",
+    "waypoints",
+    "horizontalalignment",
+    "verticalalignment",
+    "boxstyle",
+    "hpadding",
+    "vpadding",
+    "hmargin",
+    "vmargin",
+    "ports",
+)
+
+
+default = {
+    "vertex": {
+        "size": 20,
+        "facecolor": "black",
+        "marker": "o",
+        "label": {
+            "color": "white",
+            "horizontalalignment": "center",
+            "verticalalignment": "center",
+            "hpadding": 18,
+            "vpadding": 12,
+        },
+    },
+    "edge": {
+        "linewidth": 1.5,
+        "linestyle": "-",
+        "color": "black",
+        "curved": False,
+        "offset": 3,
+        "tension": 1,
+        "looptension": 4,
+        "loopmaxangle": 60,
+        "label": {
+            "horizontalalignment": "center",
+            "verticalalignment": "center",
+            "rotate": False,
+            "bbox": {
+                "boxstyle": "round",
+                "facecolor": "white",
+                "edgecolor": "none",
+            },
+        },
+        "arrow": {
+            "marker": "|>",
+            "width": 8,
+        },
+    },
+    "grouping": {
+        "facecolor": ["grey", "steelblue", "tomato"],
+        "edgecolor": "black",
+        "linewidth": 1.5,
+        "alpha": 0.5,
+        "vertexpadding": 18,
+    },
+}
+
+
+def copy_with_deep_values(style):
+    """Make a deep copy of the style dict but do not create copies of the keys."""
+    newdict = {}
+    for key, value in style.items():
+        if isinstance(value, dict):
+            newdict[key] = copy_with_deep_values(value)
+        else:
+            newdict[key] = copy.copy(value)
+    return newdict
+
+
+hollow = copy_with_deep_values(default)
+hollow["vertex"]["color"] = None
+hollow["vertex"]["facecolor"] = "none"
+hollow["vertex"]["edgecolor"] = "black"
+hollow["vertex"]["linewidth"] = 1.5
+hollow["vertex"]["marker"] = "r"
+hollow["vertex"]["size"] = "label"
+
+tree = copy_with_deep_values(default)
+tree["vertex"]["size"] = 0
+tree["vertex"]["alpha"] = 0
+tree["edge"]["linewidth"] = 2.5
+tree["vertex"]["label"]["bbox"] = {
+    "boxstyle": "square,pad=0.5",
+    "facecolor": "white",
+    "edgecolor": "none",
+}
+tree["vertex"]["label"]["color"] = "black"
+tree["vertex"]["label"]["size"] = 12
+tree["vertex"]["label"]["horizontalalignment"] = "left"
+tree["vertex"]["label"]["hmargin"] = 5
+
+
+styles = {
+    "default": default,
+    "hollow": hollow,
+    "tree": tree,
+}
+
+
+stylename = "default"
+
+
+current = copy_with_deep_values(styles["default"])
+
+
+def get_stylename():
+    """Return the name of the current iplotx style."""
+    return str(stylename)
+
+
+def get_style(name: str = "") -> dict[str, Any]:
+    """Get a *deep copy* of the chosen style.
+
+    Parameters:
+        name: The name of the style to get. If empty, the current style is returned.
+            Substyles can be obtained by using a dot notation, e.g. "default.vertex".
+            If "name" starts with a dot, it means a substyle of the current style.
+    Returns:
+        The requected style or substyle.
+
+    NOTE: The deep copy is a little different from standard deep copies. Here, keys
+        (which need to be hashables) are never copied, but values are. This can be
+        useful for hashables that change hash upon copying, such as Biopython's
+        tree nodes.
+    """
+    namelist = name.split(".")
+    style = styles
+    for i, namei in enumerate(namelist):
+        if (i == 0) and (namei == ""):
+            style = current
+        else:
+            if namei in style:
+                style = style[namei]
+            # NOTE: if asking for a nonexistent, non-leaf style
+            # give the benefit of the doubt and set an empty dict
+            # which will not fail unless the uder tries to enter it
+            elif namei not in style_leaves:
+                style = {}
+            else:
+                raise KeyError(f"Style not found: {name}")
+
+    style = copy_with_deep_values(style)
+    return style
+
+
+# The following is inspired by matplotlib's style library
+# https://github.com/matplotlib/matplotlib/blob/v3.10.3/lib/matplotlib/style/core.py#L45
+def use(style: Optional[str | dict | Sequence] = None, **kwargs):
+    """Use iplotx style setting for a style specification.
+
+    The style name of 'default' is reserved for reverting back to
+    the default style settings.
+
+    Parameters:
+        style: A style specification, currently either a name of an existing style
+            or a dict with specific parts of the style to override. The string
+            "default" resets the style to the default one. If this is a sequence,
+            each style is applied in order.
+        **kwargs: Additional style changes to be applied at the end of any style.
+    """
+    try:
+        import networkx as nx
+    except ImportError:
+        nx = None
+
+    global current
+
+    def _sanitize_leaves(style: dict):
+        for key, value in style.items():
+            if key in style_leaves:
+                if nx is not None:
+                    if isinstance(value, nx.classes.reportviews.NodeView):
+                        style[key] = dict(value)
+                    elif isinstance(value, nx.classes.reportviews.EdgeViewABC):
+                        style[key] = [v for *e, v in value]
+            elif isinstance(value, dict):
+                _sanitize_leaves(value)
+
+    def _update(style: dict, current: dict):
+        for key, value in style.items():
+            if key not in current:
+                current[key] = value
+                continue
+
+            # Style leaves are by definition not to be recurred into
+            if isinstance(value, dict) and (key not in style_leaves):
+                _update(value, current[key])
+            elif value is None:
+                del current[key]
+            else:
+                current[key] = value
+
+    old_style = copy_with_deep_values(current)
+
+    try:
+        if style is None:
+            styles = []
+        elif isinstance(style, (dict, str)):
+            styles = [style]
+        else:
+            styles = list(style)
+
+        if kwargs:
+            styles.append(kwargs)
+
+        for style in styles:
+            if style == "default":
+                reset()
+            else:
+                if isinstance(style, str):
+                    current = get_style(style)
+                else:
+                    _sanitize_leaves(style)
+                    unflatten_style(style)
+                    _update(style, current)
+    except:
+        current = old_style
+        raise
+
+
+def reset() -> Never:
+    """Reset to default style."""
+    global current
+    current = copy_with_deep_values(styles["default"])
+
+
+@contextmanager
+def context(style: Optional[str | dict | Sequence] = None, **kwargs):
+    """Create a style context for iplotx.
+
+    Parameters:
+        style: A single style specification or a list of style specifications, which are then
+            applied in order. Each style can be a string (for an existing style) or a dictionary
+            with the elements that are to change.
+        **kwargs: Additional style changes to be applied at the end of all styles.
+
+    Yields:
+        A context manager that applies the style and reverts it back to the previous one upon exit.
+    """
+    current = get_style()
+    try:
+        use(style, **kwargs)
+        yield
+    finally:
+        use(["default", current])
+
+
+def unflatten_style(
+    style_flat: dict[str, str | dict | int | float],
+) -> Never:
+    """Convert a flat or semi-flat style into a fully structured dict.
+
+    Parameters:
+        style_flat: A flat dictionary where keys may contain underscores, which are taken to signify
+            subdictionaries.
+
+    NOTE: The dict is changed *in place*.
+
+    Example:
+        >>> style = {'vertex_size': 20}
+        >>> unflatten_style(style)
+        >>> print(style)
+        {'vertex': {'size': 20}}
+    """
+
+    def _inner(style_flat: dict):
+        keys = list(style_flat.keys())
+
+        for key in keys:
+            if "_" not in key:
+                continue
+
+            keyhead, keytail = key.split("_", 1)
+            value = style_flat.pop(key)
+            if keyhead not in style_flat:
+                style_flat[keyhead] = {
+                    keytail: value,
+                }
+            else:
+                style_flat[keyhead][keytail] = value
+
+        for key, value in style_flat.items():
+            if isinstance(value, dict) and (key not in style_leaves):
+                _inner(value)
+
+    # top-level adjustments
+    if "zorder" in style_flat:
+        style_flat["network_zorder"] = style_flat["grouping_zorder"] = style_flat.pop(
+            "zorder"
+        )
+
+    # Begin recursion
+    _inner(style_flat)
+
+
+def rotate_style(
+    style: dict[str, Any],
+    index: Optional[int] = None,
+    key: Optional[Hashable] = None,
+    props: Sequence[str] = style_leaves,
+) -> dict[str, Any]:
+    """Rotate leaves of a style for a certain index or key.
+
+    Parameters:
+        style: The style to rotate.
+        index: The integer to rotate the style leaves into.
+        key: For dict-like leaves (e.g. vertex properties specified as a dict-like object over the
+            vertices themselves), the key to use for rotation (e.g. the vertex itself).
+        props: The properties to rotate, usually all leaf properties.
+
+    Returns:
+        A style with rotated leaves, which describes the properties of a single element (e.g. vertex).
+
+    Example:
+        >>> style = {'vertex': {'size': [10, 20]}}
+        >>> rotate_style(style, index=0)
+        {'vertex': {'size': 10}}
+    """
+    if (index is None) and (key is None):
+        raise ValueError(
+            "At least one of 'index' or 'key' must be provided to rotate_style."
+        )
+
+    style = copy_with_deep_values(style)
+
+    for prop in props:
+        val = style.get(prop, None)
+        if val is None:
+            continue
+        # Try integer indexing for ordered types
+        if (index is not None) and isinstance(
+            val, (tuple, list, np.ndarray, pd.Index, pd.Series)
+        ):
+            style[prop] = np.asarray(val)[index % len(val)]
+        # Try key indexing for unordered, dict-like types
+        if (
+            (key is not None)
+            and (not isinstance(val, (str, tuple, list, np.ndarray)))
+            and hasattr(val, "__getitem__")
+        ):
+            # If only a subset of keys is provided, default the other ones
+            # to the empty type constructor (e.g. 0 for ints, 0.0 for floats,
+            # empty strings).
+            if key in val:
+                style[prop] = val[key]
+            else:
+                valtype = type(next(iter(val.values())))
+                style[prop] = valtype()
+
+    return style