7b9b58429dd76c0746a64564cff0f5131f9e70a2,torch_geometric/utils/tree_decomposition.py,,tree_decomposition,#Any#,15

Before Change

    if Chem is None:
        raise ImportError("Package `rdkit` could not be found.")

    if mol.GetNumAtoms() <= 1:
        return torch.tensor((2, 0)), torch.tensor((2, 0)), 0

    // Cliques = rings and bonds.
    cliques = [list(x) for x in Chem.GetSymmSSSR(mol)]
    for bond in mol.GetBonds():
        if not bond.IsInRing():
            cliques.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()])

After Change

                    count = len(set(cliques[c1]) & set(cliques[c2]))
                    edges[(c1, c2)] = min(count, edges.get((c1, c2), 99))

    if len(edges) > 0:
        edge_index_T, weight = zip(*edges.items())
        row, col = torch.tensor(edge_index_T).t()
        inv_weight = 100 - torch.tensor(weight)
        clique_graph = SparseTensor(row=row, col=col, value=inv_weight,
                                    sparse_sizes=(len(cliques), len(cliques)))
        junc_tree = minimum_spanning_tree(clique_graph.to_scipy("csr"))
        row, col, _ = SparseTensor.from_scipy(junc_tree).coo()
        edge_index = torch.stack([row, col], dim=0)
        edge_index = to_undirected(edge_index, num_nodes=len(cliques))
    else:
        edge_index = torch.tensor((2, 0), dtype=torch.long)

    rows = [[i] * len(atom2clique[i]) for i in range(mol.GetNumAtoms())]
    row = torch.tensor(list(chain.from_iterable(rows)))
    col = torch.tensor(list(chain.from_iterable(atom2clique)))
    atom2clique = torch.stack([row, col], dim=0)

Instances