Utils

graph

Graph utilities wrapping optional NetworkX and RDFLib dependencies, plus pykeen-TriplesFactory-compatible id mapping and train/test splitting that work without pykeen installed.

build_id_mappings

build_id_mappings(
    triples: list[Triple],
) -> tuple[dict[str, int], dict[str, int]]

Build pykeen-TriplesFactory-compatible entity/relation id mappings.

Mirrors what pykeen.triples.TriplesFactory.from_labeled_triples builds internally when no explicit mapping is given: entities are the sorted union of all subjects and objects, relations are the sorted set of predicates, each enumerated from 0. Sorting makes the mapping deterministic regardless of triple order/duplication.

Build this from the complete triple set (e.g. before splitting via train_test_split_triples) so that a train split and a test split can share one consistent id space -- see that function's docstring.

Parameters:

Name	Type	Description	Default
triples	list[Triple]	Triples to derive ids from, e.g. from to_triples.	required

Returns:

Type	Description
dict[str, int]	(entity_to_id, relation_to_id). Both are {} if triples
dict[str, int]	is empty.

Source code in src/linkingtk/utils/graph.py

def build_id_mappings(triples: list[Triple]) -> tuple[dict[str, int], dict[str, int]]:
    """Build pykeen-``TriplesFactory``-compatible entity/relation id mappings.

    Mirrors what ``pykeen.triples.TriplesFactory.from_labeled_triples`` builds
    internally when no explicit mapping is given: entities are the sorted
    union of all subjects and objects, relations are the sorted set of
    predicates, each enumerated from ``0``. Sorting makes the mapping
    deterministic regardless of triple order/duplication.

    Build this from the *complete* triple set (e.g. before splitting via
    [train_test_split_triples][linkingtk.utils.graph.train_test_split_triples])
    so that a train split and a test split can share one consistent id
    space -- see that function's docstring.

    Args:
        triples: Triples to derive ids from, e.g. from
            [to_triples][linkingtk.utils.graph.to_triples].

    Returns:
        ``(entity_to_id, relation_to_id)``. Both are ``{}`` if ``triples``
        is empty.
    """
    entities = sorted({s for s, _, _ in triples} | {o for _, _, o in triples})
    relations = sorted({p for _, p, _ in triples})
    return (
        {label: index for index, label in enumerate(entities)},
        {label: index for index, label in enumerate(relations)},
    )

map_triples_to_ids

map_triples_to_ids(
    triples: list[Triple],
    entity_to_id: Mapping[str, int],
    relation_to_id: Mapping[str, int],
) -> npt.NDArray[np.int64]

Map label triples to a pykeen-TriplesFactory-ready id array.

Parameters:

Name	Type	Description	Default
triples	list[Triple]	Triples to map.	required
entity_to_id	Mapping[str, int]	Subject/object label -> id, e.g. from build_id_mappings.	required
relation_to_id	Mapping[str, int]	Predicate label -> id.	required

Returns:

Type	Description
NDArray[int64]	An (n, 3) int64 array of (head_id, relation_id, tail_id)
NDArray[int64]	rows, in the same column order pykeen's TriplesFactory.mapped_triples
NDArray[int64]	uses -- pass it straight through as that constructor's
NDArray[int64]	mapped_triples argument. Shape is (0, 3) if triples is
NDArray[int64]	empty.

Raises:

Type	Description
KeyError	If a triple references a label absent from entity_to_id/relation_to_id.

Source code in src/linkingtk/utils/graph.py

def map_triples_to_ids(
    triples: list[Triple],
    entity_to_id: Mapping[str, int],
    relation_to_id: Mapping[str, int],
) -> npt.NDArray[np.int64]:
    """Map label triples to a pykeen-``TriplesFactory``-ready id array.

    Args:
        triples: Triples to map.
        entity_to_id: Subject/object label -> id, e.g. from
            [build_id_mappings][linkingtk.utils.graph.build_id_mappings].
        relation_to_id: Predicate label -> id.

    Returns:
        An ``(n, 3)`` ``int64`` array of ``(head_id, relation_id, tail_id)``
        rows, in the same column order pykeen's ``TriplesFactory.mapped_triples``
        uses -- pass it straight through as that constructor's
        ``mapped_triples`` argument. Shape is ``(0, 3)`` if ``triples`` is
        empty.

    Raises:
        KeyError: If a triple references a label absent from
            ``entity_to_id``/``relation_to_id``.
    """
    if not triples:
        return np.empty((0, 3), dtype=np.int64)
    return np.array(
        [(entity_to_id[s], relation_to_id[p], entity_to_id[o]) for s, p, o in triples],
        dtype=np.int64,
    )

to_triples

to_triples(
    graph: Graph, *, relation_attr: str = "relation"
) -> list[Triple]

Normalize a supported graph representation into a plain triple list.

Parameters:

Name	Type	Description	Default
graph	Graph	A list of (subject, predicate, object) triples, a networkx.Graph (or DiGraph/MultiGraph/MultiDiGraph), an rdflib.Graph, or None.	required
relation_attr	str	For networkx graphs only: the edge-data key to read each triple's relation/predicate label from (e.g. set via G.add_edge(u, v, relation="author_of")). When an edge lacks this attribute, the triple falls back to the literal relation "related_to". Ignored for rdflib.Graph and plain-list inputs, which already carry real predicates.	'relation'

Returns:

Type	Description
list[Triple]	The graph as a list of triples. Empty if graph is None.

Raises:

Type	Description
TypeError	If graph is not a triple list, None, or an instance of an installed NetworkX/RDFLib graph type.

Source code in src/linkingtk/utils/graph.py

def to_triples(graph: Graph, *, relation_attr: str = "relation") -> list[Triple]:
    """Normalize a supported graph representation into a plain triple list.

    Args:
        graph: A list of ``(subject, predicate, object)`` triples, a
            ``networkx.Graph`` (or ``DiGraph``/``MultiGraph``/``MultiDiGraph``),
            an ``rdflib.Graph``, or ``None``.
        relation_attr: For ``networkx`` graphs only: the edge-data key to
            read each triple's relation/predicate label from (e.g. set via
            ``G.add_edge(u, v, relation="author_of")``). When an edge lacks
            this attribute, the triple falls back to the literal relation
            ``"related_to"``. Ignored for ``rdflib.Graph`` and plain-list
            inputs, which already carry real predicates.

    Returns:
        The graph as a list of triples. Empty if ``graph`` is ``None``.

    Raises:
        TypeError: If ``graph`` is not a triple list, ``None``, or an
            instance of an installed NetworkX/RDFLib graph type.
    """
    if graph is None:
        return []
    if isinstance(graph, list):
        return graph

    try:
        import rdflib

        if isinstance(graph, rdflib.Graph):
            return [(str(s), str(p), str(o)) for s, p, o in graph]
    except ImportError:
        pass

    try:
        import networkx as nx

        if isinstance(graph, nx.Graph):
            return [
                (str(u), str(data.get(relation_attr, "related_to")), str(v))
                for u, v, data in graph.edges(data=True)
            ]
    except ImportError:
        pass

    raise TypeError(f"Unsupported graph type: {type(graph)!r}")

train_test_split_triples

train_test_split_triples(
    triples: list[Triple],
    test_size: float = 0.2,
    random_state: int | None = None,
) -> tuple[list[Triple], list[Triple]]

Randomly split triples into train/test lists.

A thin wrapper around sklearn.model_selection.train_test_split (no pykeen dependency, so this works without the kge extra installed). Unlike pykeen.triples.TriplesFactory.split, this does not guarantee every entity/relation appears in the train split -- it's a plain uniform-random row split over triples, so an entity that happens to have all its triples land in the test partition will be absent from train (and vice versa). Callers that need that coverage guarantee should build a TriplesFactory and call its own .split() instead.

Build id mappings (via build_id_mappings) from the full triples list passed here, not from the returned splits independently, so both splits share one consistent id space -- e.g.::

train, test = train_test_split_triples(triples, test_size=0.2, random_state=0)
entity_to_id, relation_to_id = build_id_mappings(triples)
train_ids = map_triples_to_ids(train, entity_to_id, relation_to_id)
test_ids = map_triples_to_ids(test, entity_to_id, relation_to_id)

Self-loops (subject == object) and multi-edges/duplicate triples (e.g. from a networkx.MultiDiGraph via to_triples) are treated as ordinary, independent rows -- no deduplication or entity/relation-aware grouping is performed, so two triples that share an entity pair can land on opposite sides of the split.

Parameters:

Name	Type	Description	Default
triples	list[Triple]	Triples to split.	required
test_size	float	Fraction of triples assigned to the test split (sklearn's own test_size semantics).	0.2
random_state	int | None	Seed for reproducible splitting; None splits non-deterministically (sklearn's default).	None

Returns:

Type	Description
list[Triple]	(train_triples, test_triples). ([], []) if triples is
list[Triple]	empty.

Raises:

Type	Description
ValueError	Propagated from sklearn if test_size/the number of triples can't produce a non-empty train and test split (e.g. fewer than 2 triples).

Source code in src/linkingtk/utils/graph.py

def train_test_split_triples(
    triples: list[Triple],
    test_size: float = 0.2,
    random_state: int | None = None,
) -> tuple[list[Triple], list[Triple]]:
    """Randomly split triples into train/test lists.

    A thin wrapper around ``sklearn.model_selection.train_test_split`` (no
    pykeen dependency, so this works without the ``kge`` extra installed).
    Unlike ``pykeen.triples.TriplesFactory.split``, this does **not**
    guarantee every entity/relation appears in the train split -- it's a
    plain uniform-random row split over ``triples``, so an entity that
    happens to have all its triples land in the test partition will be
    absent from train (and vice versa). Callers that need that coverage
    guarantee should build a ``TriplesFactory`` and call its own
    ``.split()`` instead.

    Build id mappings (via
    [build_id_mappings][linkingtk.utils.graph.build_id_mappings]) from the
    *full* ``triples`` list passed here, not from the returned splits
    independently, so both splits share one consistent id space -- e.g.::

        train, test = train_test_split_triples(triples, test_size=0.2, random_state=0)
        entity_to_id, relation_to_id = build_id_mappings(triples)
        train_ids = map_triples_to_ids(train, entity_to_id, relation_to_id)
        test_ids = map_triples_to_ids(test, entity_to_id, relation_to_id)

    Self-loops (``subject == object``) and multi-edges/duplicate triples
    (e.g. from a ``networkx.MultiDiGraph`` via
    [to_triples][linkingtk.utils.graph.to_triples]) are treated as
    ordinary, independent rows -- no deduplication or entity/relation-aware
    grouping is performed, so two triples that share an entity pair can
    land on opposite sides of the split.

    Args:
        triples: Triples to split.
        test_size: Fraction of ``triples`` assigned to the test split
            (sklearn's own ``test_size`` semantics).
        random_state: Seed for reproducible splitting; ``None`` splits
            non-deterministically (sklearn's default).

    Returns:
        ``(train_triples, test_triples)``. ``([], [])`` if ``triples`` is
        empty.

    Raises:
        ValueError: Propagated from sklearn if ``test_size``/the number of
            triples can't produce a non-empty train and test split (e.g.
            fewer than 2 triples).
    """
    if not triples:
        return [], []
    train, test = _sk_train_test_split(triples, test_size=test_size, random_state=random_state)
    return train, test

device

Shared device-resolution helper for PyTorch-backed training code.

KGELinker is excluded -- it's pykeen-backed, and pykeen's own trainer already handles device placement. Every other linker in linkingtk.algorithms.ea (MTransELinker, IPTransELinker, JAPELinker, KDCoELinker, AttrELinker, IMUSELinker, MultiKELinker) hand-rolls its own PyTorch training loop and calls resolve_device once at the top of fit() -- as does Trainer.

resolve_device

resolve_device(device: str) -> torch.device

Resolve a device string (e.g. "cpu", "cuda", "cuda:0") to a torch.device.

Parameters:

Name	Type	Description	Default
device	str	A device string acceptable to torch.device.	required

Returns:

Type	Description
device	The resolved torch.device.

Raises:

Type	Description
LinkingTKError	If device isn't a valid device string, or requests a CUDA device while torch.cuda.is_available() is False. Deliberately doesn't fall back to CPU silently -- that would hide a real misconfiguration.

Source code in src/linkingtk/utils/device.py

def resolve_device(device: str) -> torch.device:
    """Resolve a device string (e.g. ``"cpu"``, ``"cuda"``, ``"cuda:0"``) to a `torch.device`.

    Args:
        device: A device string acceptable to `torch.device`.

    Returns:
        The resolved `torch.device`.

    Raises:
        LinkingTKError: If ``device`` isn't a valid device string, or
            requests a CUDA device while `torch.cuda.is_available()` is
            `False`. Deliberately doesn't fall back to CPU silently --
            that would hide a real misconfiguration.
    """
    import torch

    try:
        resolved = torch.device(device)
    except RuntimeError as exc:
        raise LinkingTKError(f"Invalid device {device!r}: {exc}") from exc

    if resolved.type == "cuda" and not torch.cuda.is_available():
        raise LinkingTKError(
            f"device={device!r} requested but torch.cuda.is_available() is False "
            "-- no CUDA device is visible to torch in this environment."
        )

    return resolved

sparse_gcn

Shared sparse-adjacency helpers for the GNN-based EA linkers (GCNAlignLinker, RDGCNLinker, AliNetLinker).

Every one of the three needs the same symmetric-degree-normalize step over a (weighted or unweighted) adjacency after self-loops are added -- normalize_adjacency_coo below is that one piece of math, ported once. The edge weighting itself (GCN-Align's relation-functionality weighting, RDGCN's plain unweighted primal graph, AliNet's plain unweighted 1-/2-hop graphs) is genuinely method-specific and stays in each linker's own private _<name>_training.py.

Kept as its own module (not folded into linkingtk.utils.graph) so that module can stay torch-free -- coo_to_torch_sparse imports torch lazily inside the function body, matching every other optional-torch import in this package (see resolve_device).

coo_to_torch_sparse

coo_to_torch_sparse(
    indices: NDArray[int64],
    values: NDArray[float64],
    size: tuple[int, int],
    device: device,
) -> torch.Tensor

Build a coalesced torch.sparse_coo_tensor from COO arrays.

Thin wrapper around the same torch.sparse_coo_tensor(...).coalesce() pattern already used in StructuredLogits -- reused here instead of each GNN linker's own _<name>_torch.py reimplementing it.

Parameters:

Name	Type	Description	Default
indices	NDArray[int64]	(2, nnz) int64 array, e.g. from normalize_adjacency_coo.	required
values	NDArray[float64]	(nnz,) edge weights.	required
size	tuple[int, int]	(num_rows, num_cols).	required
device	device	Target device.	required

Returns:

Type	Description
Tensor	A coalesced sparse COO tensor on device.

Source code in src/linkingtk/utils/sparse_gcn.py

def coo_to_torch_sparse(
    indices: npt.NDArray[np.int64],
    values: npt.NDArray[np.float64],
    size: tuple[int, int],
    device: torch.device,
) -> torch.Tensor:
    """Build a coalesced `torch.sparse_coo_tensor` from COO arrays.

    Thin wrapper around the same ``torch.sparse_coo_tensor(...).coalesce()``
    pattern already used in
    [StructuredLogits][linkingtk.algorithms.wsd._ewiser_structured_logits.StructuredLogits] --
    reused here instead of each GNN linker's own ``_<name>_torch.py``
    reimplementing it.

    Args:
        indices: ``(2, nnz)`` int64 array, e.g. from
            [normalize_adjacency_coo][linkingtk.utils.sparse_gcn.normalize_adjacency_coo].
        values: ``(nnz,)`` edge weights.
        size: ``(num_rows, num_cols)``.
        device: Target device.

    Returns:
        A coalesced sparse COO tensor on ``device``.
    """
    import torch

    idx = torch.from_numpy(indices).long().to(device)
    val = torch.from_numpy(values).float().to(device)
    return torch.sparse_coo_tensor(idx, val, size=size, device=device).coalesce()

normalize_adjacency_coo

normalize_adjacency_coo(
    indices: NDArray[int64],
    values: NDArray[float64],
    num_nodes: int,
    add_self_loops: bool = True,
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]

Symmetrically degree-normalize a weighted adjacency: D^-0.5 A D^-0.5.

Degree D comes from A's column sums, not row sums -- this only matters when A is asymmetric (true for GCN-Align's relation-functionality-weighted edges; a no-op for RDGCN's/AliNet's plain unweighted, inherently-symmetric adjacencies, where row- and column-sums coincide). This is a deliberately-derived match for OpenEA's own GCN_Utils.normalize_adj, not an arbitrary choice: that function computes D^-0.5 A_stored^T D^-0.5 (row-sum degree, but with an extra transpose) over a matrix A_stored that OpenEA's own get_weighted_adj already builds row/col-swapped (row.append(key[1]); col.append(key[0]) -- storing each edge at [target, source], not [source, target]). Working through both swaps algebraically: D^-0.5 A_stored^T D^-0.5 computed over that swapped storage is exactly D^-0.5 A D^-0.5 computed over the natural (source-row, target-column) orientation this function expects (see build_weighted_adjacency), using column sums for the degree (since A_stored's row-sums are A's column-sums once the swap unwinds). Confirmed empirically, not just algebraically: on the real EN-FR-15K-V1 benchmark, this column-sum version reaches Hits@10=0.243/MRR=0.154 vs. a row-sum version's Hits@10=0.148/MRR=0.116 (300 epochs, otherwise identical) -- a meaningfully different, and clearly better, result, not an inconsequential implementation detail. Duplicate (row, col) entries in indices are summed before normalizing.

Parameters:

Name	Type	Description	Default
indices	NDArray[int64]	(2, nnz) int64 array, row 0 = row ids, row 1 = column ids -- the same axis layout torch.sparse_coo_tensor expects, so coo_to_torch_sparse can consume this function's output directly.	required
values	NDArray[float64]	(nnz,) edge weights, same order as indices' columns.	required
num_nodes	int	Total node count (both KGs combined, in this package's usage -- entity ids from build_id_mappings are already in one shared id space).	required
add_self_loops	bool	Whether to add weight-1.0 self-loops before normalizing (summed with any self-loop weight already present in indices/values, not overwritten).	True

Returns:

Type	Description
NDArray[int64]	(out_indices, out_values) in the same (2, nnz)/(nnz,)
NDArray[float64]	shape as the input, coalesced (no duplicate (row, col) pairs).

Source code in src/linkingtk/utils/sparse_gcn.py

def normalize_adjacency_coo(
    indices: npt.NDArray[np.int64],
    values: npt.NDArray[np.float64],
    num_nodes: int,
    add_self_loops: bool = True,
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]:
    """Symmetrically degree-normalize a weighted adjacency: ``D^-0.5 A D^-0.5``.

    Degree ``D`` comes from ``A``'s **column** sums, not row sums -- this
    only matters when ``A`` is asymmetric (true for GCN-Align's
    relation-functionality-weighted edges; a no-op for RDGCN's/AliNet's
    plain unweighted, inherently-symmetric adjacencies, where row- and
    column-sums coincide). This is a deliberately-derived match for OpenEA's
    own ``GCN_Utils.normalize_adj``, not an arbitrary choice: that function
    computes ``D^-0.5 A_stored^T D^-0.5`` (row-sum degree, but with an
    extra transpose) over a matrix ``A_stored`` that OpenEA's own
    ``get_weighted_adj`` already builds row/col-swapped
    (``row.append(key[1]); col.append(key[0])`` -- storing each edge at
    ``[target, source]``, not ``[source, target]``). Working through both
    swaps algebraically: ``D^-0.5 A_stored^T D^-0.5`` computed over that
    swapped storage is *exactly* ``D^-0.5 A D^-0.5`` computed over the
    natural (source-row, target-column) orientation this function expects
    (see
    [build_weighted_adjacency][linkingtk.algorithms.ea._gcn_align_training.build_weighted_adjacency]),
    **using column sums for the degree** (since ``A_stored``'s row-sums are
    ``A``'s column-sums once the swap unwinds). Confirmed empirically, not
    just algebraically: on the real EN-FR-15K-V1 benchmark, this
    column-sum version reaches Hits@10=0.243/MRR=0.154 vs. a row-sum
    version's Hits@10=0.148/MRR=0.116 (300 epochs, otherwise identical) --
    a meaningfully different, and clearly better, result, not an
    inconsequential implementation detail. Duplicate ``(row, col)`` entries
    in ``indices`` are summed before normalizing.

    Args:
        indices: ``(2, nnz)`` int64 array, row 0 = row ids, row 1 = column
            ids -- the same axis layout `torch.sparse_coo_tensor` expects,
            so [coo_to_torch_sparse][linkingtk.utils.sparse_gcn.coo_to_torch_sparse]
            can consume this function's output directly.
        values: ``(nnz,)`` edge weights, same order as ``indices``' columns.
        num_nodes: Total node count (both KGs combined, in this package's
            usage -- entity ids from
            [build_id_mappings][linkingtk.utils.graph.build_id_mappings]
            are already in one shared id space).
        add_self_loops: Whether to add weight-``1.0`` self-loops before
            normalizing (summed with any self-loop weight already present
            in ``indices``/``values``, not overwritten).

    Returns:
        ``(out_indices, out_values)`` in the same ``(2, nnz)``/``(nnz,)``
        shape as the input, coalesced (no duplicate ``(row, col)`` pairs).
    """
    row, col = indices[0], indices[1]
    adjacency = sp.coo_matrix((values, (row, col)), shape=(num_nodes, num_nodes))
    adjacency.sum_duplicates()
    if add_self_loops:
        adjacency = adjacency + sp.eye(num_nodes, format="coo")
    adjacency = adjacency.tocoo()

    col_sum = np.asarray(adjacency.sum(axis=0)).flatten()
    with np.errstate(divide="ignore"):
        d_inv_sqrt = np.power(col_sum, -0.5)
    d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.0
    d_mat_inv_sqrt = sp.diags(d_inv_sqrt)

    normalized = d_mat_inv_sqrt.dot(adjacency).dot(d_mat_inv_sqrt).tocoo()
    out_indices = np.stack([normalized.row, normalized.col]).astype(np.int64)
    out_values = normalized.data.astype(np.float64)
    return out_indices, out_values