mantpy 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
mantpy/__init__.py ADDED
@@ -0,0 +1,80 @@
1
+ """
2
+ Mantpy — a framework for extracellular matrix analysis in spatial omics.
3
+
4
+ Graph-based analysis of the extracellular matrix and cell--ECM interactions in
5
+ multiplexed imaging and spatial transcriptomics.
6
+
7
+ Usage
8
+ -----
9
+ import mantpy as mt
10
+
11
+ ic = mt.im.ImageContainer("codex.tif", lazy=True) # large-image support
12
+ adata = mt.read_imc(ic, panel=panel, cells=cells)
13
+ mt.pp.extract_ecm_patches(adata, ic, tiled=True, tile_size=2048)
14
+ mt.gr.build_cell_graph(adata)
15
+ mt.gr.build_ecm_graph(adata)
16
+ mt.gr.build_cell_ecm_graph(adata)
17
+ mt.tl.interaction_test(adata)
18
+ mt.pl.cell_ecm_graph(adata)
19
+ """
20
+
21
+ from importlib import import_module as _import_module
22
+ from types import ModuleType as _ModuleType
23
+
24
+ import mantpy.datasets as datasets
25
+ import mantpy.ds as ds
26
+ import mantpy.fetch as fetch
27
+ import mantpy.gr as gr
28
+ import mantpy.im as im
29
+ import mantpy.io as io
30
+ import mantpy.palette as palette
31
+ import mantpy.pp as pp
32
+ import mantpy.style as style
33
+ import mantpy.tl as tl
34
+ from mantpy._version import __version__
35
+ from mantpy.ds import MantpyDataset
36
+ from mantpy.im import ImageContainer
37
+ from mantpy.io import read_codex, read_ecm_image, read_imc, read_imc_folder
38
+
39
+
40
+ def __getattr__(name: str) -> _ModuleType:
41
+ """Load optional and plotting namespaces only when first requested.
42
+
43
+ Plotting and optional neural dependencies can import
44
+ ``matplotlib.pyplot``, which resolves Matplotlib's backend and changes
45
+ ``rcParams["backend"]``. Keeping them lazy makes the base package import
46
+ free of plotting-state changes while preserving ``mt.pl`` and ``mt.nn``.
47
+ """
48
+ if name in {"nn", "pl"}:
49
+ module = _import_module(f"mantpy.{name}")
50
+ globals()[name] = module
51
+ return module
52
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
53
+
54
+
55
+ __all__ = [
56
+ "__version__",
57
+ "io",
58
+ "im",
59
+ "pp",
60
+ "gr",
61
+ "tl",
62
+ "pl",
63
+ "ds",
64
+ "datasets",
65
+ "fetch",
66
+ "nn",
67
+ "palette",
68
+ "style",
69
+ "read_imc",
70
+ "read_imc_folder",
71
+ "read_codex",
72
+ "read_ecm_image",
73
+ "ImageContainer",
74
+ "MantpyDataset",
75
+ ]
76
+
77
+
78
+ def __dir__() -> list[str]:
79
+ """List public lazy exports alongside the populated module namespace."""
80
+ return sorted(set(globals()) | set(__all__))
mantpy/_constants.py ADDED
@@ -0,0 +1,34 @@
1
+ CELLTYPE_COL = "cell_type"
2
+ X_COL = "x"
3
+ Y_COL = "y"
4
+ SAMPLE_ID_COL = "sample_id"
5
+ CONDITION_COL = "condition"
6
+
7
+ SPATIAL_KEY = "spatial"
8
+ RAW_LAYER = "raw"
9
+
10
+ CELL_GRAPH_KEY = "cell_graph"
11
+ ECM_GRAPH_KEY = "ecm_graph"
12
+ CELL_ECM_GRAPH_KEY = "cell_ecm_graph"
13
+
14
+ #: squidpy's canonical connectivity slot (``squidpy._constants._pkg_constants
15
+ #: .Key.obsp.spatial_conn``). Every ``sq.gr.*`` function defaults to it, so a
16
+ #: graph published here is visible to the whole squidpy toolchain.
17
+ SQUIDPY_CONN_KEY = "spatial_connectivities"
18
+
19
+ ECM_PATCHES_KEY = "ecm_patches"
20
+ ECM_CLUSTER_COL = "ecm_cluster"
21
+ ECM_IMAGE_KEY = "ecm_image"
22
+ IMAGE_CONTAINER_KEY = "image_container"
23
+
24
+ INTERACTION_TEST_KEY = "interaction_test"
25
+ NEIGHBOURHOOD_CLUSTERS_KEY = "neighbourhood_clusters"
26
+
27
+ MANTPY_UNS_KEY = "mantpy"
28
+
29
+ NODE_TYPE_CELL = "cell"
30
+ NODE_TYPE_ECM = "ecm"
31
+
32
+ EDGE_TYPE_CC = "cell-cell"
33
+ EDGE_TYPE_EE = "ecm-ecm"
34
+ EDGE_TYPE_CE = "cell-ecm"
@@ -0,0 +1 @@
1
+ # Internal modules — not part of the public API.
mantpy/_core/_edges.py ADDED
@@ -0,0 +1,549 @@
1
+ """Edge construction algorithms for spatial graphs.
2
+
3
+ Internal module — not part of the public API.
4
+
5
+ All edge-builder functions return a list of (i, j, distance) tuples so callers
6
+ can decide how to weight or threshold them.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from collections.abc import Callable
12
+
13
+ import networkx as nx
14
+ import numpy as np
15
+
16
+ EdgeList = list[tuple[int, int, float]]
17
+ WeightFn = Callable[[int, int, float], float]
18
+
19
+
20
+ def knn_edges(
21
+ coords: np.ndarray,
22
+ k: int,
23
+ Dmax: float | None = None,
24
+ ) -> EdgeList:
25
+ """K-nearest-neighbour edges.
26
+
27
+ Parameters
28
+ ----------
29
+ coords
30
+ (N, 2) array of (x, y) coordinates.
31
+ k
32
+ Number of neighbours per node.
33
+ Dmax
34
+ Maximum allowed distance. Edges longer than this are pruned.
35
+
36
+ Returns
37
+ -------
38
+ List of (i, j, distance) tuples (undirected — each edge appears once).
39
+ """
40
+ from sklearn.neighbors import NearestNeighbors
41
+
42
+ n = len(coords)
43
+ if n < 2:
44
+ return []
45
+
46
+ k_actual = min(k + 1, n) # +1 because NearestNeighbors includes self
47
+ nn = NearestNeighbors(n_neighbors=k_actual, algorithm="ball_tree")
48
+ nn.fit(coords)
49
+ distances, indices = nn.kneighbors(coords)
50
+
51
+ edges: EdgeList = []
52
+ seen: set[tuple[int, int]] = set()
53
+ for i in range(n):
54
+ for j_pos in range(1, len(indices[i])):
55
+ j = int(indices[i][j_pos])
56
+ d = float(distances[i][j_pos])
57
+ if Dmax is not None and d > Dmax:
58
+ continue
59
+ key = (min(i, j), max(i, j))
60
+ if key not in seen:
61
+ seen.add(key)
62
+ edges.append((i, j, d))
63
+ return edges
64
+
65
+
66
+ def delaunay_edges(
67
+ coords: np.ndarray,
68
+ Dmax: float | None = None,
69
+ ) -> EdgeList:
70
+ """Delaunay triangulation edges.
71
+
72
+ Parameters
73
+ ----------
74
+ coords
75
+ (N, 2) array of (x, y) coordinates.
76
+ Dmax
77
+ Maximum allowed distance.
78
+
79
+ Returns
80
+ -------
81
+ List of (i, j, distance) tuples.
82
+ """
83
+ from scipy.spatial import Delaunay
84
+
85
+ n = len(coords)
86
+ if n < 3:
87
+ return knn_edges(coords, k=min(2, n - 1), Dmax=Dmax)
88
+
89
+ tri = Delaunay(coords)
90
+ edges: EdgeList = []
91
+ seen: set[tuple[int, int]] = set()
92
+ for simplex in tri.simplices:
93
+ for a, b in [(0, 1), (1, 2), (0, 2)]:
94
+ i, j = int(simplex[a]), int(simplex[b])
95
+ key = (min(i, j), max(i, j))
96
+ if key in seen:
97
+ continue
98
+ seen.add(key)
99
+ d = float(np.linalg.norm(coords[i] - coords[j]))
100
+ if Dmax is not None and d > Dmax:
101
+ continue
102
+ edges.append((i, j, d))
103
+ return edges
104
+
105
+
106
+ def radius_edges(
107
+ coords: np.ndarray,
108
+ r: float,
109
+ ) -> EdgeList:
110
+ """Radius-based edges: connect all pairs within distance ``r``.
111
+
112
+ Parameters
113
+ ----------
114
+ coords
115
+ (N, 2) array of (x, y) coordinates.
116
+ r
117
+ Radius threshold.
118
+
119
+ Returns
120
+ -------
121
+ List of (i, j, distance) tuples.
122
+ """
123
+ from sklearn.neighbors import BallTree
124
+
125
+ n = len(coords)
126
+ if n < 2:
127
+ return []
128
+
129
+ tree = BallTree(coords)
130
+ indices = tree.query_radius(coords, r=r)
131
+ edges: EdgeList = []
132
+ seen: set[tuple[int, int]] = set()
133
+ for i, neighbours in enumerate(indices):
134
+ for j in neighbours:
135
+ j = int(j)
136
+ if j == i:
137
+ continue
138
+ key = (min(i, j), max(i, j))
139
+ if key in seen:
140
+ continue
141
+ seen.add(key)
142
+ d = float(np.linalg.norm(coords[i] - coords[j]))
143
+ edges.append((i, j, d))
144
+ return edges
145
+
146
+
147
+ def grid_edges(
148
+ coords: np.ndarray,
149
+ connectivity: int = 8,
150
+ spacing: float | None = None,
151
+ ) -> EdgeList:
152
+ """Lattice-connectivity edges for points sampled on a regular grid.
153
+
154
+ Designed for ECM patch centroids, which
155
+ :func:`~mantpy.pp.extract_ecm_patches` places on a regular, non-overlapping
156
+ grid. Each node connects to its immediate lattice neighbours — either the
157
+ 4-connected von Neumann neighbourhood (orthogonal only) or the 8-connected
158
+ Moore neighbourhood (orthogonal + diagonal). The grid spacing is inferred
159
+ from the data, so the construction is **parameter-free**: there is no
160
+ distance threshold to tune or defend.
161
+
162
+ Unlike ``knn_edges(coords, k=8)``, this never produces long edges that reach
163
+ across gaps in the lattice (e.g. where background patches were dropped) — a
164
+ node on the edge of a hole simply ends up with fewer neighbours. It is the
165
+ clean equivalent of "k=8 KNN then cap diagonals", but expressed directly as
166
+ a grid stencil.
167
+
168
+ Parameters
169
+ ----------
170
+ coords
171
+ (N, 2) array of (x, y) coordinates on a regular grid.
172
+ connectivity
173
+ ``8`` — Moore neighbourhood (4 orthogonal + 4 diagonal, default).
174
+ ``4`` — von Neumann neighbourhood (orthogonal only).
175
+ spacing
176
+ Grid spacing. ``None`` (default) infers it as the median
177
+ nearest-neighbour distance, which equals the lattice step for a regular
178
+ grid and is robust to a few duplicate/short distances.
179
+
180
+ Returns
181
+ -------
182
+ List of (i, j, distance) tuples (undirected — each edge appears once).
183
+ """
184
+ from sklearn.neighbors import BallTree
185
+
186
+ if connectivity not in (4, 8):
187
+ raise ValueError(f"connectivity must be 4 or 8, got {connectivity!r}.")
188
+
189
+ n = len(coords)
190
+ if n < 2:
191
+ return []
192
+
193
+ tree = BallTree(coords)
194
+ if spacing is None:
195
+ nn_dist, _ = tree.query(coords, k=2)
196
+ step = nn_dist[:, 1]
197
+ step = step[step > 0]
198
+ if len(step) == 0:
199
+ return []
200
+ spacing = float(np.median(step))
201
+
202
+ # Pick a query radius strictly between the farthest in-stencil neighbour and
203
+ # the nearest out-of-stencil one, so only true lattice neighbours connect:
204
+ # 4-conn: orthogonal (=spacing) < r < diagonal (=spacing*sqrt2)
205
+ # 8-conn: diagonal (=spacing*sqrt2) < r < 2-step orthogonal (=2*spacing)
206
+ sqrt2 = 2.0**0.5
207
+ radius = spacing * (2.0 + sqrt2) / 2.0 if connectivity == 8 else spacing * (1.0 + sqrt2) / 2.0
208
+
209
+ indices = tree.query_radius(coords, r=radius)
210
+ edges: EdgeList = []
211
+ seen: set[tuple[int, int]] = set()
212
+ for i, neighbours in enumerate(indices):
213
+ for j in neighbours:
214
+ j = int(j)
215
+ if j == i:
216
+ continue
217
+ key = (min(i, j), max(i, j))
218
+ if key in seen:
219
+ continue
220
+ seen.add(key)
221
+ d = float(np.linalg.norm(coords[i] - coords[j]))
222
+ edges.append((i, j, d))
223
+ return edges
224
+
225
+
226
+ def build_cell_nx_graph(
227
+ coords: np.ndarray,
228
+ cell_types: np.ndarray,
229
+ edges: EdgeList,
230
+ extra_node_attrs: dict[str, np.ndarray] | None = None,
231
+ weight_fn: WeightFn | None = None,
232
+ edge_feat_fn: WeightFn | None = None,
233
+ ) -> nx.Graph:
234
+ """Build a NetworkX graph from cell coordinates and edges.
235
+
236
+ Parameters
237
+ ----------
238
+ coords
239
+ (N, 2) array of (x, y) coordinates.
240
+ cell_types
241
+ (N,) array of cell-type strings.
242
+ edges
243
+ List of (i, j, distance) tuples.
244
+ extra_node_attrs
245
+ Optional dict of ``{attr_name: array_of_length_N}`` for additional
246
+ node attributes (e.g. marker expression).
247
+ weight_fn
248
+ Callable ``(i, j, distance) -> float`` for edge weight.
249
+ Defaults to ``1/distance``.
250
+ edge_feat_fn
251
+ Optional callable ``(i, j, distance) -> 1-D float32 array`` for edge
252
+ feature vectors. When provided, each edge stores ``feat_fwd`` (i→j)
253
+ and ``feat_rev`` (j→i) attributes, which ``to_pyg`` exports as
254
+ ``edge_attr``.
255
+
256
+ Returns
257
+ -------
258
+ nx.Graph with node attributes: ``node_type``, ``x``, ``y``, ``cell_type``,
259
+ and edge attributes: ``distance``, ``weight``, ``edge_type``.
260
+ """
261
+ from mantpy._constants import EDGE_TYPE_CC, NODE_TYPE_CELL
262
+
263
+ if weight_fn is None:
264
+ weight_fn = lambda i, j, d: 1.0 / d if d > 0 else float("inf")
265
+
266
+ G = nx.Graph()
267
+ for i in range(len(coords)):
268
+ attrs = {
269
+ "node_type": NODE_TYPE_CELL,
270
+ "x": float(coords[i, 0]),
271
+ "y": float(coords[i, 1]),
272
+ "cell_type": str(cell_types[i]),
273
+ }
274
+ if extra_node_attrs:
275
+ for k, v in extra_node_attrs.items():
276
+ attrs[k] = v[i]
277
+ G.add_node(f"cell_{i}", **attrs)
278
+
279
+ for i, j, d in edges:
280
+ w = weight_fn(i, j, d)
281
+ eattrs: dict = {"distance": d, "weight": w, "edge_type": EDGE_TYPE_CC}
282
+ if edge_feat_fn is not None:
283
+ eattrs["feat_fwd"] = edge_feat_fn(i, j, d)
284
+ eattrs["feat_rev"] = edge_feat_fn(j, i, d)
285
+ G.add_edge(f"cell_{i}", f"cell_{j}", **eattrs)
286
+
287
+ return G
288
+
289
+
290
+ def build_ecm_nx_graph(
291
+ ecm_patches_df,
292
+ edges: EdgeList,
293
+ extra_node_attrs: dict[str, list] | None = None,
294
+ weight_fn: WeightFn | None = None,
295
+ edge_feat_fn: WeightFn | None = None,
296
+ ) -> nx.Graph:
297
+ """Build a NetworkX graph from ECM patch nodes and edges.
298
+
299
+ Parameters
300
+ ----------
301
+ ecm_patches_df
302
+ DataFrame with ``x``, ``y``, ``ecm_cluster``, and feature columns.
303
+ edges
304
+ List of (i, j, distance) tuples.
305
+ extra_node_attrs
306
+ Optional additional node attributes.
307
+ weight_fn
308
+ Callable ``(i, j, distance) -> float`` for edge weight.
309
+ Defaults to ``1/distance``.
310
+ edge_feat_fn
311
+ Optional callable ``(i, j, distance) -> 1-D float32 array``.
312
+ When provided, stores ``feat_fwd`` (i→j) and ``feat_rev`` (j→i) on each
313
+ edge, exported as ``edge_attr`` by ``to_pyg``.
314
+
315
+ Returns
316
+ -------
317
+ nx.Graph with node attributes: ``node_type``, ``x``, ``y``,
318
+ ``ecm_cluster``, ``ecm_features``, and any extra attrs.
319
+ """
320
+ from mantpy._constants import EDGE_TYPE_EE, NODE_TYPE_ECM
321
+
322
+ if weight_fn is None:
323
+ weight_fn = lambda i, j, d: 1.0 / d if d > 0 else float("inf")
324
+
325
+ feat_cols = [c for c in ecm_patches_df.columns if c.startswith("feat_")]
326
+ G = nx.Graph()
327
+
328
+ for i, row in enumerate(ecm_patches_df.itertuples(index=False)):
329
+ features = [getattr(row, c) for c in feat_cols]
330
+ attrs = {
331
+ "node_type": NODE_TYPE_ECM,
332
+ "x": float(row.x),
333
+ "y": float(row.y),
334
+ "ecm_cluster": int(row.ecm_cluster),
335
+ "ecm_features": features,
336
+ }
337
+ if extra_node_attrs:
338
+ for k, v in extra_node_attrs.items():
339
+ attrs[k] = v[i]
340
+ G.add_node(f"ecm_{i}", **attrs)
341
+
342
+ for i, j, d in edges:
343
+ w = weight_fn(i, j, d)
344
+ eattrs: dict = {"distance": d, "weight": w, "edge_type": EDGE_TYPE_EE}
345
+ if edge_feat_fn is not None:
346
+ eattrs["feat_fwd"] = edge_feat_fn(i, j, d)
347
+ eattrs["feat_rev"] = edge_feat_fn(j, i, d)
348
+ G.add_edge(f"ecm_{i}", f"ecm_{j}", **eattrs)
349
+
350
+ return G
351
+
352
+
353
+ def merge_cell_ecm_graphs(
354
+ cell_graph: nx.Graph,
355
+ ecm_graph: nx.Graph,
356
+ cell_coords: np.ndarray,
357
+ ecm_coords: np.ndarray,
358
+ Dmax_CE: float | None,
359
+ k: int,
360
+ weight_fn: WeightFn | None = None,
361
+ edge_feat_fn=None,
362
+ edge_method: str = "knn",
363
+ ) -> nx.Graph:
364
+ """Merge cell and ECM graphs, adding cross edges within ``Dmax_CE``.
365
+
366
+ Parameters
367
+ ----------
368
+ cell_graph
369
+ NetworkX graph of cell nodes.
370
+ ecm_graph
371
+ NetworkX graph of ECM nodes.
372
+ cell_coords
373
+ (N_cells, 2) coordinates.
374
+ ecm_coords
375
+ (N_ecm, 2) coordinates.
376
+ Dmax_CE
377
+ Maximum cell-ECM distance to add a cross edge. ``None`` = no limit.
378
+ k
379
+ Maximum number of ECM neighbours per cell.
380
+ weight_fn
381
+ Callable ``(i, j, distance) -> float`` for cross-edge weight.
382
+ edge_feat_fn
383
+ Optional callable ``(ci, ei, distance, *, reverse=False) -> 1-D
384
+ float32 array`` for cross-edge features. ``ci`` indexes
385
+ ``cell_coords`` and ``ei`` indexes ``ecm_coords``; the callable is
386
+ invoked twice per cross edge with ``reverse=False`` for the
387
+ cell→ECM orientation and ``reverse=True`` for ECM→cell. Custom
388
+ callables MUST accept the ``reverse`` keyword — passing one that
389
+ doesn't will raise ``TypeError`` (mantpy's own
390
+ :func:`mantpy.gr.build_cell_ecm_graph` already wraps user-supplied
391
+ edge features into a closure that supports it).
392
+
393
+ Returns
394
+ -------
395
+ nx.Graph with all cell, ECM, and cross edges.
396
+ """
397
+ from mantpy._constants import EDGE_TYPE_CE
398
+
399
+ if weight_fn is None:
400
+ weight_fn = lambda i, j, d: 1.0 / d if d > 0 else float("inf")
401
+
402
+ G = nx.Graph()
403
+ G.add_nodes_from(cell_graph.nodes(data=True))
404
+ G.add_edges_from(cell_graph.edges(data=True))
405
+ G.add_nodes_from(ecm_graph.nodes(data=True))
406
+ G.add_edges_from(ecm_graph.edges(data=True))
407
+
408
+ if len(cell_coords) == 0 or len(ecm_coords) == 0:
409
+ return G
410
+
411
+ if edge_method == "delaunay":
412
+ cross_edges = delaunay_edges_bipartite(cell_coords, ecm_coords, Dmax=Dmax_CE)
413
+ elif edge_method == "knn":
414
+ cross_edges = knn_edges_bipartite(cell_coords, ecm_coords, k=k, Dmax=Dmax_CE)
415
+ else:
416
+ raise ValueError(f"edge_method must be 'knn' or 'delaunay', got {edge_method!r}")
417
+
418
+ # Node identity is recovered by name prefix, so the incoming graphs must use
419
+ # mantpy's own naming. A graph built elsewhere (squidpy, or any networkx
420
+ # default) has integer node names, which match neither prefix — see below.
421
+ cell_nodes = [n for n in G.nodes if str(n).startswith("cell_")]
422
+ ecm_nodes = [n for n in G.nodes if str(n).startswith("ecm_")]
423
+
424
+ # Cross-edges are indexed into these lists positionally. If a list is short,
425
+ # the index is unresolvable: previously we minted an attribute-less
426
+ # `cell_{ci}`/`ecm_{ei}` placeholder, which silently routed every cross-edge
427
+ # onto phantom nodes and left the real ones with none — leaving downstream
428
+ # callers to fail far away with "No cell-ECM edges with signal clusters
429
+ # found." Fail here instead, where the cause is visible.
430
+ for name, nodes, coords in (
431
+ ("cell", cell_nodes, cell_coords),
432
+ ("ecm", ecm_nodes, ecm_coords),
433
+ ):
434
+ if len(nodes) < len(coords):
435
+ raise ValueError(
436
+ f"merge_cell_ecm_graphs found {len(nodes)} '{name}_*' nodes but was given "
437
+ f"{len(coords)} {name} coordinates. Node names must be '{name}_<i>' matching "
438
+ f"the row order of {name}_coords; graphs built outside mantpy (e.g. by "
439
+ "sq.gr.spatial_neighbors) use integer names and must be relabelled first — "
440
+ f"e.g. nx.relabel_nodes(g, {{i: f'{name}_{{i}}' for i in range(g.number_of_nodes())}})."
441
+ )
442
+
443
+ for ci, ei, d in cross_edges:
444
+ cn = cell_nodes[ci]
445
+ en = ecm_nodes[ei]
446
+ w = weight_fn(ci, ei, d)
447
+ eattrs: dict = {"distance": d, "weight": w, "edge_type": EDGE_TYPE_CE}
448
+ if edge_feat_fn is not None:
449
+ # Cross-edge conventions:
450
+ # forward = cell → ecm (edge_feat_fn expects ci = cell idx, ei = ecm idx)
451
+ # reverse = ecm → cell (pass reverse=True so the closure swaps which
452
+ # coords array it indexes into)
453
+ # Custom callables passed in must accept reverse=False/True;
454
+ # the previous "swap indices" fallback silently returned wrong
455
+ # values whenever Na != Nb (cell coords and ecm coords differ
456
+ # in length, which is the normal case).
457
+ eattrs["feat_fwd"] = edge_feat_fn(ci, ei, d, reverse=False)
458
+ eattrs["feat_rev"] = edge_feat_fn(ci, ei, d, reverse=True)
459
+ G.add_edge(cn, en, **eattrs)
460
+
461
+ return G
462
+
463
+
464
+ def knn_edges_bipartite(
465
+ coords_a: np.ndarray,
466
+ coords_b: np.ndarray,
467
+ k: int,
468
+ Dmax: float | None = None,
469
+ ) -> list[tuple[int, int, float]]:
470
+ """K-nearest neighbours from set A to set B (directed A→B)."""
471
+ from sklearn.neighbors import NearestNeighbors
472
+
473
+ if len(coords_a) == 0 or len(coords_b) == 0:
474
+ return []
475
+
476
+ k_actual = min(k, len(coords_b))
477
+ nn = NearestNeighbors(n_neighbors=k_actual, algorithm="ball_tree")
478
+ nn.fit(coords_b)
479
+ distances, indices = nn.kneighbors(coords_a)
480
+
481
+ edges = []
482
+ for i in range(len(coords_a)):
483
+ for pos in range(k_actual):
484
+ j = int(indices[i][pos])
485
+ d = float(distances[i][pos])
486
+ if Dmax is not None and d > Dmax:
487
+ continue
488
+ edges.append((i, j, d))
489
+ return edges
490
+
491
+
492
+ def delaunay_edges_bipartite(
493
+ coords_a: np.ndarray,
494
+ coords_b: np.ndarray,
495
+ Dmax: float | None = None,
496
+ ) -> list[tuple[int, int, float]]:
497
+ """Bipartite Delaunay edges between two point sets.
498
+
499
+ Runs Delaunay triangulation on the stacked point cloud
500
+ ``[coords_a; coords_b]`` and keeps only simplex edges that cross the
501
+ A/B boundary. Returns ``(i_a, i_b, distance)`` tuples indexed into
502
+ the two input arrays.
503
+
504
+ Parameters
505
+ ----------
506
+ coords_a
507
+ ``(Na, 2)`` coordinates for set A.
508
+ coords_b
509
+ ``(Nb, 2)`` coordinates for set B.
510
+ Dmax
511
+ Optional distance cap. Edges longer than this are pruned.
512
+ """
513
+ from scipy.spatial import Delaunay, QhullError
514
+
515
+ na = len(coords_a)
516
+ nb = len(coords_b)
517
+ if na == 0 or nb == 0:
518
+ return []
519
+ if na + nb < 3:
520
+ return knn_edges_bipartite(coords_a, coords_b, k=min(1, nb), Dmax=Dmax)
521
+
522
+ stacked = np.vstack([coords_a, coords_b])
523
+ try:
524
+ tri = Delaunay(stacked)
525
+ except QhullError:
526
+ return knn_edges_bipartite(coords_a, coords_b, k=min(5, nb), Dmax=Dmax)
527
+
528
+ seen: set[tuple[int, int]] = set()
529
+ edges: list[tuple[int, int, float]] = []
530
+ for simplex in tri.simplices:
531
+ pts = [int(p) for p in simplex]
532
+ for a, b in [(0, 1), (1, 2), (0, 2)]:
533
+ p, q = pts[a], pts[b]
534
+ # one endpoint in A (< na), the other in B (>= na)
535
+ if (p < na) == (q < na):
536
+ continue
537
+ if p < na:
538
+ ia, ib = p, q - na
539
+ else:
540
+ ia, ib = q, p - na
541
+ key = (ia, ib)
542
+ if key in seen:
543
+ continue
544
+ seen.add(key)
545
+ d = float(np.linalg.norm(coords_a[ia] - coords_b[ib]))
546
+ if Dmax is not None and d > Dmax:
547
+ continue
548
+ edges.append((ia, ib, d))
549
+ return edges