Source code for pipeline.TrackBuilding.edges2tracks

"""
"""
import scipy.sparse as sps
import numpy as np
import torch


def build_tracks_from_edges(edge_index: torch.Tensor, n_hits: int) -> torch.Tensor:
    row, col = edge_index.cpu()
    edge_attr = np.ones(row.size(0))

    sparse_edges = sps.coo_matrix(
        (edge_attr, (row.numpy(), col.numpy())), (n_hits, n_hits)
    )

    _, candidate_labels = sps.csgraph.connected_components(
        sparse_edges, directed=False, return_labels=True
    )
    return torch.from_numpy(candidate_labels).to(edge_index.device)