bsplot 0.0.2__py3-none-any.whl

bsplot/graph/layout.py

@@ -0,0 +1,201 @@
+"""
+Utilities for layout adjustments.
+
+This module provides helpers to spread nodes that lie on the same level using
+only an existing `pos` dict.
+"""
+from __future__ import annotations
+from typing import Dict, Tuple, Optional, Callable
+import numpy as np
+
+Position = Tuple[float, float]
+
+
+def scatter_nodes(
+    pos: Dict[object, Position],
+    orientation: str = "horizontal",
+    spacing: Optional[float] = None,
+    jitter: float = 0.0,
+    tol: Optional[float] = None,
+    preserve_center: bool = True,
+    level_extractor: Optional[Callable[[Position], float]] = None,
+    seed: Optional[int] = None,
+    *,
+    spacing_unit: str = "data",
+    spacing_multiplier: Optional[float] = None,
+    # Axis-specific controls (override generic when provided)
+    spacing_x: Optional[float] = None,
+    spacing_y: Optional[float] = None,
+    spacing_unit_x: Optional[str] = None,
+    spacing_unit_y: Optional[str] = None,
+    spacing_multiplier_x: Optional[float] = None,
+    spacing_multiplier_y: Optional[float] = None,
+) -> Dict[object, Position]:
+    """
+    Return a new pos where nodes sharing the same level are evenly spread along
+    the specified orientation.
+
+    Parameters:
+    - pos: mapping of node -> (x, y)
+    - orientation: 'horizontal' to spread along x for nodes with same y;
+                   'vertical' to spread along y for nodes with same x;
+                   'both' to apply both passes (horizontal then vertical) using
+                   the original positions for level grouping.
+    - spacing: distance between adjacent nodes along the scattering axis.
+               If spacing_unit='data', interpreted in data units.
+               If spacing_unit='fraction', treated as a fraction of the global
+               span along the scattering axis (0..1). If None, a robust value is
+               inferred; optionally multiplied by spacing_multiplier.
+    - jitter: optional random jitter added along the scattering axis
+              (uniform in [-jitter/2, +jitter/2]).
+    - tol: tolerance for grouping levels; if None, exact equality defines a level
+           on the grouping axis.
+    - preserve_center: if True, keep the mean coordinate (on the scattering axis)
+                       of each level unchanged; else start at 0.
+    - level_extractor: override level extraction for single-axis modes only.
+                       Ignored for orientation='both'.
+    - seed: RNG seed for deterministic jitter.
+    - spacing_unit: 'data' (default) or 'fraction' of global axis span.
+    - spacing_multiplier: optional multiplicative factor applied to inferred
+                          spacing when spacing is None.
+    - spacing_x / spacing_y: axis-specific spacing overrides for x- or y-pass.
+    - spacing_unit_x / spacing_unit_y: axis-specific units ('data'|'fraction').
+    - spacing_multiplier_x / spacing_multiplier_y: axis-specific multipliers.
+
+    Notes:
+    - Only the scattering axis coordinate is modified; the other axis is preserved.
+    - Node order within a level is based on their current coordinate along the
+      scattering axis (ascending).
+    """
+
+    if not pos:
+        return {}
+
+    # Work on copies
+    base_pos = dict(pos)
+    new_pos: Dict[object, Position] = dict(pos)
+
+    orientation = orientation.lower()
+    if orientation not in {"horizontal", "vertical", "both"}:
+        raise ValueError("orientation must be 'horizontal', 'vertical', or 'both'")
+
+    def _norm_unit(u: Optional[str]) -> str:
+        return (u or spacing_unit).lower()
+
+    unit_x = _norm_unit(spacing_unit_x)
+    unit_y = _norm_unit(spacing_unit_y)
+
+    if unit_x not in {"data", "fraction"} or unit_y not in {"data", "fraction"}:
+        raise ValueError("spacing units must be 'data' or 'fraction'")
+
+    rng = np.random.default_rng(seed) if jitter and jitter > 0 else None
+
+    # Global spans used for fraction-based spacing
+    xs = np.array([xy[0] for xy in base_pos.values()], dtype=float)
+    ys = np.array([xy[1] for xy in base_pos.values()], dtype=float)
+    span_x = float(xs.max() - xs.min()) if xs.size else 1.0
+    span_y = float(ys.max() - ys.min()) if ys.size else 1.0
+    if span_x == 0:
+        span_x = 1.0
+    if span_y == 0:
+        span_y = 1.0
+
+    def group_levels(values: Dict[object, Position], extractor: Callable[[Position], float]):
+        # Deterministic ordering by (level, x, y)
+        sorted_items = sorted(values.items(), key=lambda kv: (extractor(kv[1]), kv[1][0], kv[1][1]))
+        groups = []  # list of (rep_level_value, [nodes])
+        for node, xy in sorted_items:
+            lvl = extractor(xy)
+            placed = False
+            if tol is not None:
+                for rep, nodes in groups:
+                    if abs(lvl - rep) <= tol:
+                        nodes.append(node)
+                        placed = True
+                        break
+            else:
+                for rep, nodes in groups:
+                    if lvl == rep:
+                        nodes.append(node)
+                        placed = True
+                        break
+            if not placed:
+                groups.append((lvl, [node]))
+        return groups
+
+    def infer_spacing(coords: np.ndarray, axis_idx: int) -> float:
+        # Choose axis-specific overrides
+        if axis_idx == 0:
+            s_val = spacing_x if spacing_x is not None else spacing
+            s_unit = unit_x
+            s_mult = spacing_multiplier_x if spacing_multiplier_x is not None else spacing_multiplier
+            base_span = span_x
+        else:
+            s_val = spacing_y if spacing_y is not None else spacing
+            s_unit = unit_y
+            s_mult = spacing_multiplier_y if spacing_multiplier_y is not None else spacing_multiplier
+            base_span = span_y
+
+        # Explicit spacing
+        if s_val is not None:
+            if s_unit == "data":
+                base = float(s_val)
+            else:  # fraction
+                base = float(s_val) * base_span
+        else:
+            # Inferred spacing from coords
+            n = coords.size
+            if n >= 2:
+                diffs = np.diff(np.sort(coords))
+                diffs = diffs[diffs > 0]
+                base = float(np.median(diffs)) if diffs.size > 0 else 1.0
+            else:
+                base = 1.0
+        if s_mult is not None:
+            base *= float(s_mult)
+        return base
+
+    def apply_pass(groups, axis_idx: int):
+        for _, nodes_on_level in groups:
+            coords = np.array([new_pos[n][axis_idx] for n in nodes_on_level], dtype=float)
+            n = len(nodes_on_level)
+            level_spacing = infer_spacing(coords, axis_idx)
+            if preserve_center:
+                center = float(coords.mean()) if n > 0 else 0.0
+                start = center - level_spacing * (n - 1) / 2.0
+            else:
+                start = 0.0
+            order = np.argsort(coords)
+            ordered_nodes = [nodes_on_level[i] for i in order]
+            for i, node in enumerate(ordered_nodes):
+                x_old, y_old = new_pos[node]
+                coord_new = start + i * level_spacing
+                if rng is not None:
+                    coord_new += float(rng.uniform(-jitter / 2.0, jitter / 2.0))
+                if axis_idx == 0:
+                    new_pos[node] = (coord_new, y_old)
+                else:
+                    new_pos[node] = (x_old, coord_new)
+
+    if orientation == "horizontal":
+        extractor = (lambda xy: xy[1]) if level_extractor is None else level_extractor
+        groups_y = group_levels(new_pos, extractor)
+        apply_pass(groups_y, axis_idx=0)
+        return new_pos
+
+    if orientation == "vertical":
+        extractor = (lambda xy: xy[0]) if level_extractor is None else level_extractor
+        groups_x = group_levels(new_pos, extractor)
+        apply_pass(groups_x, axis_idx=1)
+        return new_pos
+
+    # orientation == 'both': use original base_pos for grouping to avoid interference
+    groups_y = group_levels(base_pos, lambda xy: xy[1])  # same y-levels
+    apply_pass(groups_y, axis_idx=0)
+    groups_x = group_levels(base_pos, lambda xy: xy[0])  # same x-levels
+    apply_pass(groups_x, axis_idx=1)
+    return new_pos
+
+
+# Backward-compatible alias (deprecated)
+scatter_nodes_on_same_level_horizontally = scatter_nodes