ramify 1.1.0__py3-none-any.whl

ramify/__init__.py

@@ -0,0 +1,13 @@
+from .centerline import extract, Network
+from .partition import allocate, voronoi, subdivide
+from .width import widths, region_widths
+
+__all__ = [
+    "extract",
+    "Network",
+    "allocate",
+    "voronoi",
+    "subdivide",
+    "widths",
+    "region_widths",
+]

ramify/_io.py

@@ -0,0 +1,107 @@
+# _io.py
+from dataclasses import dataclass
+import numpy as np
+
+
+@dataclass
+class GridMeta:
+    """Everything needed to reconstruct an xr.DataArray from a numpy result."""
+
+    coords: dict
+    dims: tuple
+    crs: object | None
+    transform: object | None
+
+
+def unwrap(arr, pixel_size=None):
+    """
+    Accept np.ndarray or xr.DataArray.
+    Returns (values: np.ndarray, pixel_size: float, meta: GridMeta | None).
+    """
+    try:
+        import xarray as xr
+
+        is_xr = isinstance(arr, xr.DataArray)
+    except ImportError:
+        is_xr = False
+
+    if not is_xr:
+        return np.asarray(arr), float(pixel_size or 1.0), None
+
+    meta = GridMeta(
+        coords=arr.coords,
+        dims=arr.dims,
+        crs=getattr(arr.rio, "crs", None) if hasattr(arr, "rio") else None,
+        transform=arr.rio.transform() if hasattr(arr, "rio") else None,
+    )
+    if pixel_size is None and meta.transform is not None:
+        pixel_size = abs(
+            meta.transform.a
+        )  # from affine transform, assumes square pixels
+    return arr.values, float(pixel_size or 1.0), meta
+
+
+def wrap(values, meta):
+    """np.ndarray + meta -> xr.DataArray; passthrough if meta is None."""
+    if meta is None:
+        return values
+    import xarray as xr
+
+    out = xr.DataArray(values, coords=meta.coords, dims=meta.dims)
+    if meta.crs is not None:
+        out.rio.write_crs(meta.crs, inplace=True)
+        out.rio.write_transform(meta.transform, inplace=True)
+    return out
+
+
+def region_groups(label_arr, mask_bool=None):
+    """Group a labeled raster's pixels by label, in one pass.
+
+    Yields ``(label, flat_idx)`` per positive label, ascending, where
+    ``flat_idx`` indexes the raveled grid; ``mask_bool`` optionally restricts
+    which pixels count. This replaces the ``label_arr == label`` scan a
+    per-region loop would otherwise do, which costs the whole grid once per
+    region. The sort is stable, so each group's indices come out ascending --
+    callers rely on that for ``searchsorted`` neighbour lookups and for
+    reading off a bounding box.
+    """
+    positive = label_arr > 0
+    flat = np.flatnonzero(positive if mask_bool is None else positive & mask_bool)
+    if flat.size == 0:
+        return
+    labels = label_arr.ravel()[flat]
+    order = np.argsort(labels, kind="stable")
+    flat = flat[order]
+    labels = labels[order]
+    starts = np.flatnonzero(np.r_[True, labels[1:] != labels[:-1]])
+    for lo, hi in zip(starts, np.append(starts[1:], flat.size)):
+        yield int(labels[lo]), flat[lo:hi]
+
+
+def edt_field(mask_bool, open_boundary):
+    """Foreground array for the local-half-width distance transform.
+
+    ``distance_transform_edt`` of the returned array gives, at each shape pixel,
+    the distance to the nearest *wall*. By default (``open_boundary is None``)
+    every non-shape pixel is a wall -- the original behaviour. When
+    ``open_boundary`` is given, its truthy pixels are treated as void (open,
+    not a wall): they join the shape as foreground, so half-widths are measured
+    only to the remaining real walls. Accepts np.ndarray or xr.DataArray.
+    """
+    if open_boundary is None:
+        return mask_bool
+    open_arr, _, _ = unwrap(open_boundary)
+    open_bool = np.asarray(open_arr) > 0
+    if open_bool.shape != mask_bool.shape:
+        raise ValueError(
+            f"open_boundary shape {open_bool.shape} does not match "
+            f"mask shape {mask_bool.shape}"
+        )
+    field = mask_bool | open_bool
+    if field.all():
+        import warnings
+
+        warnings.warn(
+            "open_boundary leaves no wall pixels; local half-widths will be zero"
+        )
+    return field

ramify/centerline.py

@@ -0,0 +1,356 @@
+import heapq
+from dataclasses import dataclass, field
+
+import numpy as np
+import pandas as pd
+from scipy.ndimage import distance_transform_edt
+from skimage.morphology import skeletonize
+from skimage.graph import MCP_Geometric
+
+from ._io import unwrap, wrap, edt_field, GridMeta
+
+SQRT2 = np.sqrt(2.0)
+_OFFSETS = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]
+
+Pixel = tuple[int, int]
+
+
+@dataclass
+class Network:
+    segments: pd.DataFrame  # segment_id, path_id, strahler, length, weight,
+    #                         downstream_segment_id, pixels
+    shape: tuple[int, int]
+    pixel_size: float
+    root: Pixel
+    tips: list[Pixel]
+    _meta: GridMeta | None = field(default=None, repr=False)
+
+    def rasterize(self, by=None):
+        if by is not None and by not in ("path", "segment"):
+            raise ValueError(f"by must be None, 'path', or 'segment', got {by!r}")
+
+        flat, counts = self._pixel_index()
+        if by is None:
+            arr = np.zeros(self.shape, dtype=np.uint8)
+            arr.ravel()[flat] = 1
+            return wrap(arr, self._meta)
+
+        column = "path_id" if by == "path" else "segment_id"
+        values = np.repeat(self.segments[column].to_numpy(), counts)
+
+        # A junction pixel belongs to several segments, so the writes collide.
+        # The old loop wrote paths in descending path_id and let later writes
+        # win, which means: lowest path_id takes the pixel (so the mainstem,
+        # path_id == 1, wins), and within one path the later row in `segments`
+        # takes it. Reproduce that by sorting the collisions together and
+        # keeping one winner each, rather than replaying the writes in order --
+        # same answer, and no O(paths^2) regrouping to get there.
+        path_ids = np.repeat(self.segments["path_id"].to_numpy(), counts)
+        seq = np.repeat(np.arange(len(self.segments)), counts)
+        order = np.lexsort((-seq, path_ids, flat))  # last key sorts first
+        winner = order[np.flatnonzero(
+            np.r_[True, flat[order][1:] != flat[order][:-1]]
+        )]
+
+        arr = np.zeros(self.shape, dtype=np.uint32)
+        arr.ravel()[flat[winner]] = values[winner]
+        return wrap(arr, self._meta)
+
+    def _pixel_index(self):
+        # Flat indices of every segment's pixels, concatenated in `segments`
+        # row order, plus each segment's pixel count so the per-segment columns
+        # can be np.repeat'd out to match.
+        pixels = [np.asarray(p).reshape(-1, 2) for p in self.segments["pixels"]]
+        counts = np.fromiter((len(p) for p in pixels), np.intp, len(pixels))
+        rc = np.concatenate(pixels) if pixels else np.empty((0, 2), np.intp)
+        return rc[:, 0].astype(np.intp) * self.shape[1] + rc[:, 1], counts
+
+    def to_gdf(self):
+        if self._meta is None or self._meta.transform is None:
+            raise ValueError(
+                "to_gdf requires a georeferenced xr.DataArray input to extract()"
+            )
+        import geopandas as gpd
+        import rasterio.transform
+        from shapely.geometry import LineString
+
+        keep, geoms = [], []
+        for i, pixels in enumerate(self.segments["pixels"]):
+            if len(pixels) < 2:
+                continue
+            rows, cols = zip(*pixels)
+            xs, ys = rasterio.transform.xy(self._meta.transform, rows, cols)
+            geoms.append(LineString(zip(xs, ys)))
+            keep.append(i)
+
+        gdf = self.segments.iloc[keep].drop(columns=["pixels"]).copy()
+        gdf["geometry"] = geoms
+        return gpd.GeoDataFrame(gdf, geometry="geometry", crs=self._meta.crs)
+
+
+def extract(mask, root, tips=None, path_by="area", pixel_size=None,
+            open_boundary=None) -> Network:
+    if path_by not in ("area", "length", "strahler"):
+        raise ValueError(
+            f"path_by must be 'area', 'length', or 'strahler', got {path_by!r}"
+        )
+
+    mask_arr, px, meta = unwrap(mask, pixel_size)
+    mask_bool = mask_arr == 1
+    root = (int(root[0]), int(root[1]))
+    if not mask_bool[root]:
+        raise ValueError(f"root {root} is not inside the mask")
+
+    # 1. skeletonize
+    nodes = _skeleton_nodes(mask_bool)
+    if not nodes:
+        raise ValueError("skeletonization produced no pixels")
+
+    # 2-3. trace root and provided tips onto the skeleton (mask-constrained)
+    points = [root] + ([tuple(map(int, t)) for t in tips] if tips else [])
+    traces = _snap_paths(points, nodes, mask_bool)
+    if traces[0] is None:
+        raise ValueError(f"root {root} cannot reach the skeleton within the mask")
+    snapped_tips = []
+    if tips:
+        for t, tr in zip(points[1:], traces[1:]):
+            if tr is None:
+                raise ValueError(f"tip {t} cannot reach the skeleton within the mask")
+        snapped_tips = points[1:]
+    for tr in traces:
+        nodes.update(tr)
+
+    # 4-5a. shortest-path tree from root (guarantees the result is a tree,
+    # even when mask holes create skeleton loops)
+    parent, dist = _dijkstra_tree(nodes, root)
+
+    # 5b. resolve tips
+    if tips:
+        for t in snapped_tips:
+            if t not in parent:
+                raise ValueError(f"tip {t} is not connected to the root")
+        tip_nodes = snapped_tips
+    else:
+        tip_nodes = [n for n in _endpoints(nodes) if n != root and n in parent]
+        if not tip_nodes:
+            raise ValueError(
+                "no tips found: skeleton has no endpoints reachable from root"
+            )
+
+    # 5c. keep only pixels on some tip -> root path
+    kept = set()
+    for t in tip_nodes:
+        n = t
+        while n is not None and n not in kept:
+            kept.add(n)
+            n = parent[n]
+
+    # 6. orient and break into segments at tips and junctions
+    segments = _to_segments(kept, parent, tip_nodes, root)
+
+    # 7. annotate
+    edt = distance_transform_edt(edt_field(mask_bool, open_boundary)) * px
+    df = _annotate(segments, edt, px, path_by, root)
+
+    return Network(
+        segments=df,
+        shape=mask_bool.shape,
+        pixel_size=px,
+        root=root,
+        tips=[t for t in tip_nodes],
+        _meta=meta,
+    )
+
+
+# -- skeleton and snapping ---------------------------------------------------
+
+
+def _skeleton_nodes(mask_bool) -> set:
+    skel = skeletonize(mask_bool)
+    rows, cols = np.nonzero(skel)
+    return set(zip(rows.tolist(), cols.tolist()))
+
+
+def _snap_paths(points, nodes, mask_bool):
+    # least-cost path from the skeleton to each point, constrained to the mask;
+    # returns list of pixel-paths (or None if unreachable), aligned with points
+    penalty = np.where(mask_bool, 1.0, np.inf)
+    mcp = MCP_Geometric(penalty)
+    # only points off the skeleton need tracing; pass them as ends so the flood
+    # stops once they are reached (they sit on/near the skeleton) instead of
+    # filling the whole array. Costs and tracebacks for the reached ends are
+    # identical to the full-flood result.
+    ends = [list(p) for p in points if p not in nodes]
+    if ends:
+        mcp.find_costs(starts=[list(n) for n in nodes], ends=ends)
+
+    out = []
+    for p in points:
+        if p in nodes:
+            out.append([p])
+            continue
+        try:
+            path = mcp.traceback(list(p))
+        except ValueError:
+            out.append(None)
+            continue
+        out.append([(int(r), int(c)) for r, c in path] if path else None)
+    return out
+
+
+# -- tree construction -------------------------------------------------------
+
+
+def _dijkstra_tree(nodes, root):
+    # parent[n] is the neighbor of n one step closer to root (downstream)
+    dist = {root: 0.0}
+    parent = {root: None}
+    heap = [(0.0, root)]
+    while heap:
+        d, n = heapq.heappop(heap)
+        if d > dist[n]:
+            continue
+        r, c = n
+        for dr, dc in _OFFSETS:
+            m = (r + dr, c + dc)
+            if m not in nodes:
+                continue
+            nd = d + (SQRT2 if dr and dc else 1.0)
+            if nd < dist.get(m, np.inf):
+                dist[m] = nd
+                parent[m] = n
+                heapq.heappush(heap, (nd, m))
+    return parent, dist
+
+
+def _endpoints(nodes):
+    out = []
+    for r, c in nodes:
+        degree = sum((r + dr, c + dc) in nodes for dr, dc in _OFFSETS)
+        if degree == 1:
+            out.append((r, c))
+    return out
+
+
+def _to_segments(kept, parent, tip_nodes, root):
+    # segments are ordered upstream -> downstream; junction pixels are shared:
+    # last pixel of each upstream segment, first pixel of the downstream one
+    n_children = {}
+    for n in kept:
+        p = parent[n]
+        if p is not None:
+            n_children[p] = n_children.get(p, 0) + 1
+
+    breakpoints = set(tip_nodes) | {n for n, k in n_children.items() if k > 1}
+    stops = breakpoints | {root}
+
+    segments = []
+    for s in breakpoints:
+        if s == root:
+            continue
+        seg = [s]
+        cur = parent[s]
+        while cur not in stops:
+            seg.append(cur)
+            cur = parent[cur]
+        seg.append(cur)
+        segments.append(seg)
+    return segments
+
+
+# -- annotation ---------------------------------------------------------------
+
+
+def _annotate(segments, edt, pixel_size, path_by, root):
+    n = len(segments)
+    start_of = {seg[0]: i for i, seg in enumerate(segments)}
+    downstream = [start_of.get(seg[-1]) for seg in segments]  # None at root
+    children = [[] for _ in range(n)]
+    for i, d in enumerate(downstream):
+        if d is not None:
+            children[d].append(i)
+
+    length = np.array([_seg_length(s, pixel_size) for s in segments])
+    weight = np.array([_seg_weight(s, edt, pixel_size) for s in segments])
+
+    # post-order accumulation (iterative)
+    strahler = np.zeros(n, dtype=int)
+    sub_length = np.zeros(n)
+    sub_weight = np.zeros(n)
+    outlets = [i for i, d in enumerate(downstream) if d is None]
+    order = []
+    stack = list(outlets)
+    while stack:
+        i = stack.pop()
+        order.append(i)
+        stack.extend(children[i])
+    for i in reversed(order):  # leaves first
+        if not children[i]:
+            strahler[i] = 1
+            sub_length[i] = length[i]
+            sub_weight[i] = weight[i]
+        else:
+            orders = strahler[children[i]]
+            m = orders.max()
+            strahler[i] = m + 1 if (orders == m).sum() > 1 else m
+            sub_length[i] = length[i] + sub_length[children[i]].max()
+            sub_weight[i] = weight[i] + sub_weight[children[i]].max()
+
+    key = {
+        "area": lambda i: (sub_weight[i],),
+        "length": lambda i: (sub_length[i],),
+        "strahler": lambda i: (strahler[i], sub_length[i]),
+    }[path_by]
+
+    # heavy-path decomposition: walk upstream from each outlet, continuing
+    # along the heaviest child; other children start new paths
+    path_id = np.zeros(n, dtype=int)
+    next_id = 1
+    from collections import deque
+
+    queue = deque(sorted(outlets, key=key, reverse=True))
+    while queue:
+        cur = queue.popleft()
+        if path_id[cur]:
+            continue
+        while True:
+            path_id[cur] = next_id
+            preds = [c for c in children[cur] if not path_id[c]]
+            if not preds:
+                break
+            preds.sort(key=key, reverse=True)
+            queue.extend(preds[1:])
+            cur = preds[0]
+        next_id += 1
+
+    df = pd.DataFrame(
+        {
+            "segment_id": np.arange(1, n + 1),
+            "path_id": path_id,
+            "strahler": strahler,
+            "length": length,
+            "weight": weight,
+            "downstream_segment_id": pd.array(
+                [d + 1 if d is not None else pd.NA for d in downstream],
+                dtype="Int64",
+            ),
+            "pixels": segments,
+        }
+    )
+    return df.sort_values(["path_id", "segment_id"], ignore_index=True)
+
+
+def _seg_length(pixels, pixel_size):
+    total = 0.0
+    for (r1, c1), (r2, c2) in zip(pixels[:-1], pixels[1:]):
+        total += SQRT2 if (r1 != r2 and c1 != c2) else 1.0
+    return total * pixel_size
+
+
+def _seg_weight(pixels, edt, pixel_size):
+    # trapezoid rule for integral of distance-to-edge along the segment (~ area/2)
+    total = 0.0
+    for (r1, c1), (r2, c2) in zip(pixels[:-1], pixels[1:]):
+        step = SQRT2 if (r1 != r2 and c1 != c2) else 1.0
+        total += 0.5 * (edt[r1, c1] + edt[r2, c2]) * step
+    return total * pixel_size

ramify/partition.py

@@ -0,0 +1,215 @@
+import numpy as np
+from scipy.ndimage import distance_transform_edt
+from scipy.sparse import csr_matrix
+from scipy.sparse.csgraph import dijkstra
+from skimage.segmentation import watershed
+
+from ._io import unwrap, wrap, edt_field, region_groups
+from .centerline import Network
+
+SQRT2 = np.sqrt(2.0)
+# forward-only neighbour offsets; directed=False makes each bidirectional
+_EDGES = [(0, 1, 1.0), (1, 0, 1.0), (1, 1, SQRT2), (1, -1, SQRT2)]
+
+
+def allocate(mask, seeds, open_boundary=None, progress=None):
+    # Ordered, radius-limited claiming. Each path (seed label) claims the mask
+    # pixels within *some* of its seeds' local half-width, measured as a
+    # boundary-respecting (geodesic) distance -- so a wide-but-farther seed can
+    # reach a pixel a narrow-but-nearer seed cannot. Paths are processed
+    # biggest-first (label 1 = mainstem = highest priority) and a pixel is kept
+    # by the first (biggest) path to reach it, so wide branches claim
+    # proportionally more space at junctions. Unclaimed remainder is
+    # watershed-filled so the labels cover the mask completely.
+    #
+    # Each path's claim is a windowed, radius-bounded Dijkstra (see _reach), so
+    # cost is O(tube area) per path rather than O(domain); the mask never
+    # changes during the loop, so there is no per-step graph rebuild.
+    #
+    # `progress`, if given, is called once per path just before its Dijkstra:
+    # progress(i, n_paths, label, window_pixel_count). The window area is the
+    # honest cost proxy -- a path's seed count is its *length*, which says little
+    # about the tube it will claim. Paths run biggest-first, so the early ones are
+    # by far the slowest; a bar weighted by window area tracks that, one counting
+    # paths does not.
+    mask_arr, _, meta = unwrap(mask)
+    seed_arr, _, _ = unwrap(seeds)
+    mask_bool = mask_arr == 1
+    _check(mask_bool, seed_arr)
+
+    H, W = mask_bool.shape
+    radius = distance_transform_edt(edt_field(mask_bool, open_boundary))  # local half-width
+    allocation = np.zeros(mask_bool.shape, dtype=np.uint32)
+
+    # group seed pixels by label once, so each path works from its own coords
+    # (and bbox) instead of scanning the full seed array every iteration
+    flat = np.flatnonzero(seed_arr)
+    labels_flat = seed_arr.ravel()[flat]
+    order = np.argsort(labels_flat, kind="stable")
+    flat = flat[order]
+    labels_sorted = labels_flat[order]
+    uniq, starts = np.unique(labels_sorted, return_index=True)  # ascending
+    bounds = np.append(starts, labels_sorted.size)
+
+    for i, label in enumerate(uniq):  # ascending label == biggest path first
+        idx = flat[bounds[i]:bounds[i + 1]]
+        rr, cc = idx // W, idx % W
+        keep = mask_bool[rr, cc]
+        if not keep.any():
+            continue
+        rr, cc = rr[keep], cc[keep]
+        rad = radius[rr, cc]
+        R = float(rad.max())  # farthest this path can reach = search bound
+
+        pad = int(np.ceil(R)) + 1
+        r0, r1 = max(int(rr.min()) - pad, 0), min(int(rr.max()) + pad + 1, H)
+        c0, c1 = max(int(cc.min()) - pad, 0), min(int(cc.max()) + pad + 1, W)
+
+        # reported here, not at the top of the loop: the window is the first
+        # point where this path's cost is actually known
+        if progress is not None:
+            progress(i, len(uniq), int(label), (r1 - r0) * (c1 - c0))
+
+        seed_local = np.stack([rr - r0, cc - c0], axis=1)
+        tube = _reach(mask_bool[r0:r1, c0:c1], seed_local, rad, R)
+        sub = allocation[r0:r1, c0:c1]
+        sub[tube & (sub == 0)] = label  # keep only where no bigger path won
+
+    claimed = allocation > 0
+    unclaimed = mask_bool & ~claimed
+    if unclaimed.any() and claimed.any():
+        allocation = watershed(
+            image=distance_transform_edt(~claimed),
+            markers=allocation,
+            mask=mask_bool,
+        ).astype(np.uint32)
+
+    return wrap(allocation, meta)
+
+
+def _reach(mask_win, seed_rc, radii, R):
+    # Pixels within some seed's local half-width, geodesically, inside the
+    # window. A virtual super-source is wired to each seed s with edge weight
+    # R - radius(s) (>= 0), so one bounded search from it gives, per pixel q,
+    #   dist_V(q) = R + min_s ( geodist(q, s) - radius(s) ),
+    # and dist_V(q) <= R is exactly "some seed's radius reaches q". The bound
+    # (limit=R) keeps the search inside the tube. Metric is octile ({1, sqrt2}),
+    # matching centerline._dijkstra_tree.
+    h, w = mask_win.shape
+    ys, xs = np.nonzero(mask_win)
+    M = ys.size
+    if M == 0:
+        return np.zeros((h, w), dtype=bool)
+    ids = np.full((h, w), -1, dtype=np.int64)  # pixel -> node id, -1 off-mask
+    ids[ys, xs] = np.arange(M)
+    V = M  # super-source node id
+
+    rows, cols, wts = [], [], []
+    for dr, dc, step in _EDGES:  # grid edges between adjacent in-mask pixels
+        ny, nx = ys + dr, xs + dc
+        ok = (ny >= 0) & (ny < h) & (nx >= 0) & (nx < w)
+        nb = np.where(ok, ids[np.clip(ny, 0, h - 1), np.clip(nx, 0, w - 1)], -1)
+        keep = nb >= 0
+        rows.append(ids[ys[keep], xs[keep]])
+        cols.append(nb[keep])
+        wts.append(np.full(int(keep.sum()), step))
+
+    sids = ids[seed_rc[:, 0], seed_rc[:, 1]]  # super-source -> each seed
+    ok = sids >= 0
+    rows.append(np.full(int(ok.sum()), V))
+    cols.append(sids[ok])
+    wts.append(R - radii[ok])
+
+    n = M + 1
+    graph = csr_matrix(
+        (np.concatenate(wts), (np.concatenate(rows), np.concatenate(cols))),
+        shape=(n, n),
+    )
+    dist = dijkstra(graph, directed=False, indices=V, limit=R)
+    out = np.zeros((h, w), dtype=bool)
+    out[ys, xs] = dist[:M] <= R
+    return out
+
+
+def voronoi(mask, seeds):
+    # Nearest-seed partition of the mask: every pixel goes to the seed label
+    # it can reach by the shortest within-mask route. No ordering, no radius
+    # limits. Use for simple subdivision, e.g. splitting a path's territory
+    # by segment: voronoi(regions == path_id, segment_seeds).
+    mask_arr, _, meta = unwrap(mask)
+    seed_arr, _, _ = unwrap(seeds)
+    mask_bool = mask_arr == 1
+    _check(mask_bool, seed_arr)
+
+    markers = np.where(mask_bool, seed_arr, 0).astype(np.int64)
+    if not (markers > 0).any():
+        return wrap(np.zeros(mask_bool.shape, dtype=np.uint32), meta)
+
+    out = watershed(
+        image=distance_transform_edt(markers == 0),
+        markers=markers,
+        mask=mask_bool,
+    ).astype(np.uint32)
+    return wrap(out, meta)
+
+
+def subdivide(regions, network: Network):
+    # Subdivide each path's territory (from allocate) into segment-level
+    # territories. Each territory is seeded only by its own path's segments,
+    # so neighboring paths' labels (e.g. shared junction pixels) never bleed
+    # across boundaries. Pixels in territories whose path has no segments in
+    # the table remain 0.
+    reg_arr, _, meta = unwrap(regions)
+    if reg_arr.shape != network.shape:
+        raise ValueError(
+            f"regions shape {reg_arr.shape} does not match network grid {network.shape}"
+        )
+
+    # Group the territories once, then work each path inside its own bounding
+    # box. Everything below is local: a full-grid pass per path would cost the
+    # whole raster ~once per path, and voronoi() runs a distance transform and
+    # a watershed, so that is the expensive kind of pass. Cropping is exact
+    # here -- the distance transform measures to the nearest seed and every one
+    # of this path's seeds is inside its own bbox, and the watershed only ever
+    # floods within the territory.
+    territories = dict(region_groups(reg_arr))
+
+    out = np.zeros(network.shape, dtype=np.uint32)
+    _, w = network.shape
+    for path_id, group in network.segments.groupby("path_id"):
+        idx = territories.get(int(path_id))
+        if idx is None:
+            continue  # path swallowed during allocation
+        rows, cols = np.divmod(idx, w)
+        r0, r1 = int(rows.min()), int(rows.max()) + 1
+        c0, c1 = int(cols.min()), int(cols.max()) + 1
+
+        territory = np.zeros((r1 - r0, c1 - c0), dtype=np.uint8)
+        territory[rows - r0, cols - c0] = 1
+
+        seeds = np.zeros(territory.shape, dtype=np.uint32)
+        for _, row in group.iterrows():
+            rc = np.asarray(row["pixels"])
+            inside = (
+                (rc[:, 0] >= r0) & (rc[:, 0] < r1) & (rc[:, 1] >= c0) & (rc[:, 1] < c1)
+            )
+            rc = rc[inside]
+            seeds[rc[:, 0] - r0, rc[:, 1] - c0] = row["segment_id"]
+        seeds = np.where(territory == 1, seeds, 0)
+        if not (seeds > 0).any():
+            continue
+
+        sub = np.asarray(voronoi(territory, seeds))
+        hit = sub > 0
+        out[r0:r1, c0:c1][hit] = sub[hit]
+
+    return wrap(out, meta)
+
+
+def _check(mask_bool, seed_arr):
+    if seed_arr.shape != mask_bool.shape:
+        raise ValueError(
+            f"seeds shape {seed_arr.shape} does not match mask shape {mask_bool.shape}"
+        )
+    if (seed_arr[mask_bool] > 0).sum() == 0:
+        raise ValueError("no seed pixels found inside the mask")

ramify/width.py

@@ -0,0 +1,264 @@
+import warnings
+
+import numpy as np
+from scipy.ndimage import distance_transform_edt
+from scipy.sparse import csr_matrix
+from scipy.sparse.csgraph import connected_components
+from scipy.sparse.linalg import cg
+from scipy.spatial import cKDTree
+
+from ._io import unwrap, wrap, edt_field, region_groups
+
+SQRT2 = np.sqrt(2.0)
+# 8-connected stencil; diagonals weighted 1/sqrt(2) for isotropy (and so a pixel
+# attached only diagonally is never isolated)
+_OFFSETS = [
+    (-1, 0, 1.0),
+    (1, 0, 1.0),
+    (0, -1, 1.0),
+    (0, 1, 1.0),
+    (-1, -1, 1 / SQRT2),
+    (-1, 1, 1 / SQRT2),
+    (1, -1, 1 / SQRT2),
+    (1, 1, 1 / SQRT2),
+]
+
+
+def widths(mask, centerline, method="laplace", pixel_size=None, open_boundary=None):
+    # Per-pixel width of the shape. Exact widths (2 * distance-to-boundary)
+    # are taken at centerline pixels and interpolated across the mask.
+    # method="laplace": smooth diffusion (Laplace equation, Dirichlet BCs at
+    #   the centerline) — continuous fields, best for downstream analysis.
+    # method="nearest": each pixel takes the width of its nearest centerline
+    #   pixel (a Voronoi-style assignment, cf. ramify.voronoi) — piecewise
+    #   constant, fast, exact at the centerline.
+    if method not in ("laplace", "nearest"):
+        raise ValueError(f"method must be 'laplace' or 'nearest', got {method!r}")
+
+    mask_arr, px, meta = unwrap(mask, pixel_size)
+    cl_arr, _, _ = unwrap(centerline)
+    mask_bool = mask_arr == 1
+    cl_bool = (cl_arr > 0) & mask_bool
+    if cl_arr.shape != mask_bool.shape:
+        raise ValueError(
+            f"centerline shape {cl_arr.shape} does not match mask shape {mask_bool.shape}"
+        )
+    if not cl_bool.any():
+        raise ValueError("no centerline pixels found inside the mask")
+
+    seed_widths = np.where(
+        cl_bool, distance_transform_edt(edt_field(mask_bool, open_boundary)) * px * 2.0, 0.0
+    )
+
+    out = np.full(mask_bool.shape, np.nan)
+    if method == "nearest":
+        out[mask_bool] = _nearest(cl_bool, seed_widths)[mask_bool]
+    else:
+        idx = np.flatnonzero(mask_bool)
+        out.ravel()[idx] = _laplace(
+            idx, cl_bool.ravel()[idx], seed_widths.ravel()[idx], mask_bool.shape
+        )
+        # parts of the mask the centerline cannot reach (a detached blob, an
+        # island) have no Dirichlet data at all -- fall back rather than
+        # reporting the zero the homogeneous solve would give
+        leftover = mask_bool & np.isnan(out)
+        if leftover.any():
+            warnings.warn(
+                f"{int(leftover.sum())} mask pixels are not connected to the "
+                "centerline; filled by nearest centerline width"
+            )
+            out[leftover] = _nearest(cl_bool, seed_widths)[leftover]
+
+    out = wrap(out, meta)
+    if meta is not None:
+        try:
+            out.rio.write_nodata(np.nan, inplace=True)
+        except AttributeError:
+            pass
+    return out
+
+
+def region_widths(mask, centerline, regions, method="laplace", pixel_size=None,
+                  open_boundary=None, progress=None):
+    # Like widths(), but interpolated independently within each labeled
+    # region (e.g. the output of ramify.allocate), so widths do not diffuse
+    # across path boundaries at junctions. Each region is seeded only by the
+    # centerline pixels inside it. Regions containing no centerline pixels
+    # are filled by nearest-centerline fallback (with a warning) so the
+    # output always covers the mask.
+    #
+    # Regions are solved from their own flat pixel indices (grouped once, up
+    # front) rather than by re-scanning the full grid per label, so the cost is
+    # O(mask area) in total instead of O(mask area x number of regions).
+    #
+    # `progress`, if given, is called once per region just before that region is
+    # solved: progress(i, n_regions, label, region_pixel_count). Regions are very
+    # uneven and the laplace solve costs ~O(n^1.5) in a region's pixel count, so
+    # drive a bar off that count -- counting regions makes it race to ~99% and
+    # then sit on the few big ones for most of the wall time.
+    if method not in ("laplace", "nearest"):
+        raise ValueError(f"method must be 'laplace' or 'nearest', got {method!r}")
+
+    mask_arr, px, meta = unwrap(mask, pixel_size)
+    cl_arr, _, _ = unwrap(centerline)
+    reg_arr, _, _ = unwrap(regions)
+    mask_bool = mask_arr == 1
+    cl_bool = (cl_arr > 0) & mask_bool
+    if cl_arr.shape != mask_bool.shape or reg_arr.shape != mask_bool.shape:
+        raise ValueError("mask, centerline, and regions must share one shape")
+    if not cl_bool.any():
+        raise ValueError("no centerline pixels found inside the mask")
+
+    seed_widths = np.where(
+        cl_bool, distance_transform_edt(edt_field(mask_bool, open_boundary)) * px * 2.0, 0.0
+    )
+    cl_flat = cl_bool.ravel()
+    seed_flat = seed_widths.ravel()
+
+    out = np.full(mask_bool.shape, np.nan)
+    out_flat = out.ravel()
+    groups = list(region_groups(reg_arr, mask_bool))  # views, so the total is free
+    for i, (label, idx) in enumerate(groups):
+        if progress is not None:
+            progress(i, len(groups), label, idx.size)
+        is_seed = cl_flat[idx]
+        if not is_seed.any():
+            continue  # filled by fallback below
+        seed_vals = seed_flat[idx]
+        if method == "laplace":
+            out_flat[idx] = _laplace(idx, is_seed, seed_vals, mask_bool.shape)
+        else:
+            out_flat[idx] = _nearest_flat(idx, is_seed, seed_vals, mask_bool.shape)
+
+    leftover = mask_bool & np.isnan(out)
+    if leftover.any():
+        warnings.warn(
+            f"{int(leftover.sum())} mask pixels fall in regions with no "
+            "centerline pixels (or outside any region); filled by nearest "
+            "centerline width"
+        )
+        fallback = _nearest(cl_bool, seed_widths)
+        out[leftover] = fallback[leftover]
+
+    out = np.where(mask_bool, out, np.nan)
+    out = wrap(out, meta)
+    if meta is not None:
+        try:
+            out.rio.write_nodata(np.nan, inplace=True)
+        except AttributeError:
+            pass
+    return out
+
+
+def _neighbours(qidx, idx, dy, dx, shape):
+    # Membership test for one stencil offset: for each query pixel qidx[i], does
+    # its neighbour at (dy, dx) belong to the sorted set `idx`? Works purely on
+    # flat indices, so nothing the size of the grid (or even of the region's
+    # bounding box) is allocated. Returns (has, pos) with idx[pos[i]] the
+    # neighbour wherever has[i]. The explicit column check is what stops a
+    # dx = -1 step at column 0 from wrapping onto the previous row.
+    h, w = shape
+    rows, cols = np.divmod(qidx, w)
+    nr, nc = rows + dy, cols + dx
+    ok = (nr >= 0) & (nr < h) & (nc >= 0) & (nc < w)
+    nflat = qidx + (dy * w + dx)
+    pos = np.searchsorted(idx, nflat)
+    np.clip(pos, 0, idx.size - 1, out=pos)
+    return ok & (idx[pos] == nflat), pos
+
+
+def _nearest(cl_bool, seed_widths):
+    # Whole-grid nearest-centerline assignment (used by widths() and as the
+    # region_widths fallback); the EDT is linear in grid size.
+    _, idx = distance_transform_edt(~cl_bool, return_indices=True)
+    return seed_widths[idx[0], idx[1]]
+
+
+def _nearest_flat(idx, is_seed, seed_vals, shape):
+    # Same rule restricted to one region's seeds. A KD-tree over the region's
+    # centerline pixels costs O(region size); an EDT here would cost O(grid).
+    w = shape[1]
+    seed_rc = np.column_stack(np.divmod(idx[is_seed], w))
+    all_rc = np.column_stack(np.divmod(idx, w))
+    _, nn = cKDTree(seed_rc).query(all_rc, k=1)
+    return seed_vals[is_seed][nn]
+
+
+def _laplace(idx, is_seed, seed_vals, shape):
+    # Laplace interpolation over the pixel set `idx` (sorted flat indices) with
+    # Dirichlet BCs at the seeds. The seed rows are eliminated rather than
+    # carried as identity rows, so the system solved is
+    #
+    #     (D - W_ff) x_f = W_fs s
+    #
+    # over the free pixels only: symmetric, diagonally dominant, positive
+    # definite — which is what cg actually requires, and smaller besides.
+    free = ~is_seed
+    n_free = int(free.sum())
+    out = seed_vals.copy()
+    if n_free == 0:
+        return out
+
+    fidx = idx[free]
+    fpos = np.cumsum(free) - 1  # position in `idx` -> row in the free system
+    local = np.arange(n_free)
+
+    rows, cols, data = [], [], []
+    diag = np.zeros(n_free)
+    b = np.zeros(n_free)
+    touches_seed = np.zeros(n_free, dtype=bool)
+    for dy, dx, wt in _OFFSETS:
+        has, pos = _neighbours(fidx, idx, dy, dx, shape)
+        p = pos[has]
+        diag[has] += wt
+        nbr_seed = is_seed[p]
+        rows.append(local[has][~nbr_seed])
+        cols.append(fpos[p[~nbr_seed]])
+        data.append(np.full(int((~nbr_seed).sum()), -wt))
+        # each free pixel meets a given offset at most once, but summing over
+        # the eight offsets still accumulates, so go through bincount
+        b += np.bincount(
+            local[has][nbr_seed],
+            weights=wt * seed_vals[p[nbr_seed]],
+            minlength=n_free,
+        )
+        touches_seed[local[has][nbr_seed]] = True
+
+    rows.append(local)
+    cols.append(local)
+    data.append(diag)
+
+    A = csr_matrix(
+        (np.concatenate(data), (np.concatenate(rows), np.concatenate(cols))),
+        shape=(n_free, n_free),
+    )
+
+    x = _solve(A, b)
+
+    # A pixel set can fall apart into chunks that no seed touches (a region
+    # split by a junction, an island). Such a chunk is a singular Neumann block
+    # with a zero rhs, and because A is block diagonal cg leaves it at exactly
+    # 0.0 -- so the solver silently reports width 0 there. Exact zero is also
+    # the only way a *seeded* block can land on 0 (the maximum principle bounds
+    # it below by its smallest seed width), which makes it a cheap filter;
+    # confirm against the connectivity, then hand the chunk back as NaN for the
+    # caller to fill by fallback.
+    if (x == 0.0).any():
+        n_comp, comp = connected_components(A, directed=False)
+        seeded = np.zeros(n_comp, dtype=bool)
+        seeded[comp[touches_seed]] = True
+        x[~seeded[comp]] = np.nan
+
+    out[free] = x
+    return out
+
+
+def _solve(A, b):
+    # rtol bounds the residual, not the error, and the two diverge as a region
+    # grows (A's smallest eigenvalue shrinks): at rtol=1e-4 a 340k-pixel region
+    # lands ~1-2 width units off the exact solution. 1e-6 costs ~40% more
+    # iterations and pulls that back to ~0.03.
+    x, info = cg(A, b, rtol=1e-6)
+    if info != 0:
+        warnings.warn("conjugate gradient solver did not converge")
+    return x

ramify-1.1.0.dist-info/METADATA

@@ -0,0 +1,183 @@
+Metadata-Version: 2.4
+Name: ramify
+Version: 1.1.0
+Summary: Segment and measure binary branching shapes: centerline extraction, hierarchical space partitioning, and width interpolation
+License-File: LICENSE
+Keywords: centerline,geomorphology,medial-axis,morphometry,raster,river,segmentation,skeleton,valley,width
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: GIS
+Requires-Python: >=3.12
+Requires-Dist: geopandas
+Requires-Dist: networkx
+Requires-Dist: numpy
+Requires-Dist: rasterio
+Requires-Dist: rioxarray
+Requires-Dist: scikit-image
+Requires-Dist: scipy
+Requires-Dist: shapely
+Requires-Dist: xarray
+Provides-Extra: dev
+Requires-Dist: matplotlib; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# ramify
+
+Characterize binary branching shapes (e.g. rivers, floodplains, glaciers,
+roots, veins...). Given a shape mask, a root point, and (optionally) branch
+tips, `ramify` extracts a topology-aware centerline network, decomposes it into
+hierarchically ordered paths, allocates every pixel of the shape to its path,
+and estimates local width everywhere.
+
+![graphical abstract](https://raw.githubusercontent.com/avkoehl/ramify/main/assets/abstract.png)
+
+## Install
+
+```bash
+pip install ramify
+```
+
+Development (clone, then sync with dev extras):
+
+```bash
+git clone https://github.com/avkoehl/ramify.git
+cd ramify
+uv sync --extra dev
+```
+
+## Usage
+
+```python
+import ramify
+from ramify.data import load
+
+mask, root, tips = load()                       # bundled toy dataset
+
+net = ramify.extract(mask, root, tips=tips)     # centerline network of ordered paths
+regions = ramify.allocate(mask, net.rasterize(by="path"))
+widths = ramify.region_widths(mask, net.rasterize(), regions)
+
+net.segments                                    # DataFrame: segment_id, path_id, strahler,
+                                                #   length, weight, downstream_segment_id
+regions                                         # labeled raster: each pixel -> its path
+widths                                          # float raster: local width everywhere
+```
+
+Inputs are `np.ndarray` (with `pixel_size=`) or georeferenced `xr.DataArray`;
+outputs match the input type. `root` and `tips` are `(row, col)` pixel coordinates.
+
+## Components
+
+Each individual component is presented below.
+
+### Centerlines
+
+```python
+net = ramify.extract(mask, root, tips=tips)
+```
+
+Skeletonizes the mask, routes from each tip to the root (pruning everything else),
+and decomposes the network into ordered paths — `path_id == 1` is the mainstem.
+
+
+![extract with tips](https://raw.githubusercontent.com/avkoehl/ramify/main/assets/extract_tips.png)
+
+```python
+net = ramify.extract(mask, root)
+```
+
+Without tips, every skeleton endpoint becomes a tip.
+
+![extract auto tips](https://raw.githubusercontent.com/avkoehl/ramify/main/assets/extract_auto.png)
+
+Tips and root can often be derived automatically — glacier branch tips
+[Kienholz et al.,
+2014](https://tc.copernicus.org/articles/8/503/2014/tc-8-503-2014.pdf), channel
+initiation points, or the lowest point on the boundary as the root — or simply
+digitized in GIS software.
+
+
+### Partitioning
+
+```python
+regions = ramify.allocate(mask, net.rasterize(by="path"))
+```
+
+Assigns every pixel to a path: paths claim territory in priority order, each limited
+by the local shape radius, so wide branches claim proportionally more space at junctions.
+
+![allocate](https://raw.githubusercontent.com/avkoehl/ramify/main/assets/allocate.png)
+
+```python
+regions = ramify.voronoi(mask, net.rasterize(by="path"))
+```
+
+Nearest-centerline partition — no ordering, no radius limits.
+
+![voronoi](https://raw.githubusercontent.com/avkoehl/ramify/main/assets/voronoi.png)
+
+```python
+seg_regions = ramify.subdivide(regions, net)
+```
+
+Subdivides each path's territory further: within a territory, every pixel goes to
+its nearest centerline segment of that same path.
+
+![subdivide](https://raw.githubusercontent.com/avkoehl/ramify/main/assets/subdivide.png)
+
+### Widths
+
+Exact widths (twice the distance to the boundary) are taken at the centerline and
+interpolated across the shape. That interpolation runs either over the whole shape
+or independently within each region, which keeps junction-zone pixels from
+averaging between a branch and its mainstem:
+
+```python
+w = ramify.widths(mask, net.rasterize())
+w = ramify.region_widths(mask, net.rasterize(), regions)
+```
+
+![widths domain](https://raw.githubusercontent.com/avkoehl/ramify/main/assets/widths_domain.png)
+
+Either call also takes `method="nearest"`, which gives each pixel the width of its
+nearest centerline pixel instead of diffusing smoothly from it — piecewise
+constant, and much faster:
+
+```python
+w = ramify.widths(mask, net.rasterize(), method="nearest")
+```
+
+![widths nearest](https://raw.githubusercontent.com/avkoehl/ramify/main/assets/widths_nearest.png)
+
+## Open boundaries
+
+Everything above measures local half-width as the distance from each pixel to the
+shape's boundary, and that half-width drives three things: which branch is the
+mainstem, how far each path claims territory, and the width field. By default
+every boundary pixel is treated as a **wall**. Sometimes part of the boundary is
+not a real wall — the shape is truncated by open water, the data extent, or
+another medium — and treating it as one makes the half-width collapse to zero
+there.
+
+Pass `open_boundary`: a binary mask, on the same grid as the shape, marking the
+non-wall (void) pixels. Distances are then measured only to the remaining real
+walls. It is optional — omitted, every boundary is a wall (the behaviour above) —
+and accepted by `extract`, `allocate`, `widths`, and `region_widths`. Give it to
+every step, so all three stages measure against the same walls:
+
+```python
+net = ramify.extract(mask, root, tips=tips, open_boundary=open_boundary)
+regions = ramify.allocate(mask, net.rasterize(by="path"), open_boundary=open_boundary)
+widths = ramify.region_widths(mask, net.rasterize(), regions,
+                              open_boundary=open_boundary)
+```
+
+Below, the same mask, root, and tips are reused, but the void past the outlet is
+marked open (shaded red), and the outlet `widths` no longer taper to the cut edge.
+Mark a region with depth rather than a thin skin along the boundary: distances are
+measured *through* the open void, so a one-pixel rind would only push the wall out
+by one pixel. 
+
+![open boundary](https://raw.githubusercontent.com/avkoehl/ramify/main/assets/open_boundary.png)
+
+Only the widths are shown because on this shape the partitioning didn't change.

ramify-1.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+ramify/__init__.py,sha256=B5QnMUkVoqKTgz8oYj44B2l8UnzbQ84u3RYEgrVkcdE,262
+ramify/_io.py,sha256=PzSWuvL8-r0tIQ_qDNejhVnkJbujRFTC6qq5mBiHrEU,3656
+ramify/centerline.py,sha256=GSVKY5G2okB6HP60b_AxoYJI1bP__tE1igTnigfa3XY,12350
+ramify/partition.py,sha256=m7lELiuEEc6-R9V8NMqqoEeJYWbkN9TNHkn8PpGxtQ8,9013
+ramify/width.py,sha256=JlCMvWs35QhfOfX1dgQ-mOsB5rWypEyNXsLp7soH2Eg,10586
+ramify-1.1.0.dist-info/METADATA,sha256=phKJ00tF5KsG3h8gBCYA7bv4c0Xiy7733gCxyhcWtgA,6649
+ramify-1.1.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+ramify-1.1.0.dist-info/licenses/LICENSE,sha256=9OJ32gzj_uZUUO657WiSLpRj7FlcLoj4yZKBiE8V0kA,1069
+ramify-1.1.0.dist-info/RECORD,,

ramify-1.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.31.0
+Root-Is-Purelib: true
+Tag: py3-none-any

ramify-1.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Arthur Koehl
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.