parkererickson/LGGM-Text2Graph

1import torch2import utils3 4 5class DummyExtraFeatures:6    def __init__(self):7        """ This class does not compute anything, just returns empty tensors."""8 9    def __call__(self, noisy_data):10        X = noisy_data['X_t']11        E = noisy_data['E_t']12        y = noisy_data['y_t']13        empty_x = X.new_zeros((*X.shape[:-1], 0))14        empty_e = E.new_zeros((*E.shape[:-1], 0))15        empty_y = y.new_zeros((y.shape[0], 0))16        return utils.PlaceHolder(X=empty_x, E=empty_e, y=empty_y)17 18 19class ExtraFeatures:20    def __init__(self, extra_features_type, max_n_nodes):21        self.max_n_nodes = max_n_nodes22        self.ncycles = NodeCycleFeatures()23        self.features_type = extra_features_type24        if extra_features_type in ['eigenvalues', 'all']:25            self.eigenfeatures = EigenFeatures(mode=extra_features_type)26 27    def __call__(self, noisy_data):28        n = noisy_data['node_mask'].sum(dim=1).unsqueeze(1) / self.max_n_nodes29        x_cycles, y_cycles = self.ncycles(noisy_data)       # (bs, n_cycles)30 31        if self.features_type == 'cycles':32            E = noisy_data['E_t']33            extra_edge_attr = torch.zeros((*E.shape[:-1], 0)).type_as(E)34            return utils.PlaceHolder(X=x_cycles, E=extra_edge_attr, y=torch.hstack((n, y_cycles)))35 36        elif self.features_type == 'eigenvalues':37            eigenfeatures = self.eigenfeatures(noisy_data)38            E = noisy_data['E_t']39            extra_edge_attr = torch.zeros((*E.shape[:-1], 0)).type_as(E)40            n_components, batched_eigenvalues = eigenfeatures   # (bs, 1), (bs, 10)41            return utils.PlaceHolder(X=x_cycles, E=extra_edge_attr, y=torch.hstack((n, y_cycles, n_components,42                                                                                    batched_eigenvalues)))43        elif self.features_type == 'all':44            eigenfeatures = self.eigenfeatures(noisy_data)45            E = noisy_data['E_t']46            extra_edge_attr = torch.zeros((*E.shape[:-1], 0)).type_as(E)47            n_components, batched_eigenvalues, nonlcc_indicator, k_lowest_eigvec = eigenfeatures   # (bs, 1), (bs, 10),48                                                                                                # (bs, n, 1), (bs, n, 2)49 50            return utils.PlaceHolder(X=torch.cat((x_cycles, nonlcc_indicator, k_lowest_eigvec), dim=-1),51                                     E=extra_edge_attr,52                                     y=torch.hstack((n, y_cycles, n_components, batched_eigenvalues)))53        else:54            raise ValueError(f"Features type {self.features_type} not implemented")55 56 57class NodeCycleFeatures:58    def __init__(self):59        self.kcycles = KNodeCycles()60 61    def __call__(self, noisy_data):62        adj_matrix = noisy_data['E_t'][..., 1:].sum(dim=-1).float()63 64        x_cycles, y_cycles = self.kcycles.k_cycles(adj_matrix=adj_matrix)   # (bs, n_cycles)65        x_cycles = x_cycles.type_as(adj_matrix) * noisy_data['node_mask'].unsqueeze(-1)66        # Avoid large values when the graph is dense67        x_cycles = x_cycles / 1068        y_cycles = y_cycles / 1069        x_cycles[x_cycles > 1] = 170        y_cycles[y_cycles > 1] = 171        return x_cycles, y_cycles72 73 74class EigenFeatures:75    """76    Code taken from : https://github.com/Saro00/DGN/blob/master/models/pytorch/eigen_agg.py77    """78    def __init__(self, mode):79        """ mode: 'eigenvalues' or 'all' """80        self.mode = mode81 82    def __call__(self, noisy_data):83        E_t = noisy_data['E_t']84        mask = noisy_data['node_mask']85        A = E_t[..., 1:].sum(dim=-1).float() * mask.unsqueeze(1) * mask.unsqueeze(2)86        L = compute_laplacian(A, normalize=False)87        mask_diag = 2 * L.shape[-1] * torch.eye(A.shape[-1]).type_as(L).unsqueeze(0)88        mask_diag = mask_diag * (~mask.unsqueeze(1)) * (~mask.unsqueeze(2))89        L = L * mask.unsqueeze(1) * mask.unsqueeze(2) + mask_diag90 91        if self.mode == 'eigenvalues':92            eigvals = torch.linalg.eigvalsh(L)        # bs, n93            eigvals = eigvals.type_as(A) / torch.sum(mask, dim=1, keepdim=True)94 95            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)96            return n_connected_comp.type_as(A), batch_eigenvalues.type_as(A)97 98        elif self.mode == 'all':99            eigvals, eigvectors = torch.linalg.eigh(L)100            eigvals = eigvals.type_as(A) / torch.sum(mask, dim=1, keepdim=True)101            eigvectors = eigvectors * mask.unsqueeze(2) * mask.unsqueeze(1)102            # Retrieve eigenvalues features103            n_connected_comp, batch_eigenvalues = get_eigenvalues_features(eigenvalues=eigvals)104 105            # Retrieve eigenvectors features106            nonlcc_indicator, k_lowest_eigenvector = get_eigenvectors_features(vectors=eigvectors,107                                                                               node_mask=noisy_data['node_mask'],108                                                                               n_connected=n_connected_comp)109            return n_connected_comp, batch_eigenvalues, nonlcc_indicator, k_lowest_eigenvector110        else:111            raise NotImplementedError(f"Mode {self.mode} is not implemented")112 113 114def compute_laplacian(adjacency, normalize: bool):115    """116    adjacency : batched adjacency matrix (bs, n, n)117    normalize: can be None, 'sym' or 'rw' for the combinatorial, symmetric normalized or random walk Laplacians118    Return:119        L (n x n ndarray): combinatorial or symmetric normalized Laplacian.120    """121    diag = torch.sum(adjacency, dim=-1)     # (bs, n)122    n = diag.shape[-1]123    D = torch.diag_embed(diag)      # Degree matrix      # (bs, n, n)124    combinatorial = D - adjacency                        # (bs, n, n)125 126    if not normalize:127        return (combinatorial + combinatorial.transpose(1, 2)) / 2128 129    diag0 = diag.clone()130    diag[diag == 0] = 1e-12131 132    diag_norm = 1 / torch.sqrt(diag)            # (bs, n)133    D_norm = torch.diag_embed(diag_norm)        # (bs, n, n)134    L = torch.eye(n).unsqueeze(0) - D_norm @ adjacency @ D_norm135    L[diag0 == 0] = 0136    return (L + L.transpose(1, 2)) / 2137 138 139def get_eigenvalues_features(eigenvalues, k=5):140    """141    values : eigenvalues -- (bs, n)142    node_mask: (bs, n)143    k: num of non zero eigenvalues to keep144    """145    ev = eigenvalues146    bs, n = ev.shape147    n_connected_components = (ev < 1e-5).sum(dim=-1)148    # assert (n_connected_components > 0).all(), (n_connected_components, ev)149 150    to_extend = max(n_connected_components) + k - n151    if to_extend > 0:152        eigenvalues = torch.hstack((eigenvalues, 2 * torch.ones(bs, to_extend).type_as(eigenvalues)))153    indices = torch.arange(k).type_as(eigenvalues).long().unsqueeze(0) + n_connected_components.unsqueeze(1)154    first_k_ev = torch.gather(eigenvalues, dim=1, index=indices)155    return n_connected_components.unsqueeze(-1), first_k_ev156 157 158def get_eigenvectors_features(vectors, node_mask, n_connected, k=2):159    """160    vectors (bs, n, n) : eigenvectors of Laplacian IN COLUMNS161    returns:162        not_lcc_indicator : indicator vectors of largest connected component (lcc) for each graph  -- (bs, n, 1)163        k_lowest_eigvec : k first eigenvectors for the largest connected component   -- (bs, n, k)164    """165    bs, n = vectors.size(0), vectors.size(1)166 167    # Create an indicator for the nodes outside the largest connected components168    first_ev = torch.round(vectors[:, :, 0], decimals=3) * node_mask                        # bs, n169    # Add random value to the mask to prevent 0 from becoming the mode170    random = torch.randn(bs, n, device=node_mask.device) * (~node_mask)                                   # bs, n171    first_ev = first_ev + random172    most_common = torch.mode(first_ev, dim=1).values                                    # values: bs -- indices: bs173    mask = ~ (first_ev == most_common.unsqueeze(1))174    not_lcc_indicator = (mask * node_mask).unsqueeze(-1).float()175 176    # Get the eigenvectors corresponding to the first nonzero eigenvalues177    to_extend = max(n_connected) + k - n178    if to_extend > 0:179        vectors = torch.cat((vectors, torch.zeros(bs, n, to_extend).type_as(vectors)), dim=2)   # bs, n , n + to_extend180    indices = torch.arange(k).type_as(vectors).long().unsqueeze(0).unsqueeze(0) + n_connected.unsqueeze(2)    # bs, 1, k181    indices = indices.expand(-1, n, -1)                                               # bs, n, k182    first_k_ev = torch.gather(vectors, dim=2, index=indices)       # bs, n, k183    first_k_ev = first_k_ev * node_mask.unsqueeze(2)184 185    return not_lcc_indicator, first_k_ev186 187def batch_trace(X):188    """189    Expect a matrix of shape B N N, returns the trace in shape B190    :param X:191    :return:192    """193    diag = torch.diagonal(X, dim1=-2, dim2=-1)194    trace = diag.sum(dim=-1)195    return trace196 197 198def batch_diagonal(X):199    """200    Extracts the diagonal from the last two dims of a tensor201    :param X:202    :return:203    """204    return torch.diagonal(X, dim1=-2, dim2=-1)205 206 207class KNodeCycles:208    """ Builds cycle counts for each node in a graph.209    """210 211    def __init__(self):212        super().__init__()213 214    def calculate_kpowers(self):215        self.k1_matrix = self.adj_matrix.float()216        self.d = self.adj_matrix.sum(dim=-1)217        self.k2_matrix = self.k1_matrix @ self.adj_matrix.float()218        self.k3_matrix = self.k2_matrix @ self.adj_matrix.float()219        self.k4_matrix = self.k3_matrix @ self.adj_matrix.float()220        self.k5_matrix = self.k4_matrix @ self.adj_matrix.float()221        self.k6_matrix = self.k5_matrix @ self.adj_matrix.float()222 223    def k3_cycle(self):224        """ tr(A ** 3). """225        c3 = batch_diagonal(self.k3_matrix)226        return (c3 / 2).unsqueeze(-1).float(), (torch.sum(c3, dim=-1) / 6).unsqueeze(-1).float()227 228    def k4_cycle(self):229        diag_a4 = batch_diagonal(self.k4_matrix)230        c4 = diag_a4 - self.d * (self.d - 1) - (self.adj_matrix @ self.d.unsqueeze(-1)).sum(dim=-1)231        return (c4 / 2).unsqueeze(-1).float(), (torch.sum(c4, dim=-1) / 8).unsqueeze(-1).float()232 233    def k5_cycle(self):234        diag_a5 = batch_diagonal(self.k5_matrix)235        triangles = batch_diagonal(self.k3_matrix)236        c5 = diag_a5 - 2 * triangles * self.d - (self.adj_matrix @ triangles.unsqueeze(-1)).sum(dim=-1) + triangles237        return (c5 / 2).unsqueeze(-1).float(), (c5.sum(dim=-1) / 10).unsqueeze(-1).float()238 239    def k6_cycle(self):240        term_1_t = batch_trace(self.k6_matrix)241        term_2_t = batch_trace(self.k3_matrix ** 2)242        term3_t = torch.sum(self.adj_matrix * self.k2_matrix.pow(2), dim=[-2, -1])243        d_t4 = batch_diagonal(self.k2_matrix)244        a_4_t = batch_diagonal(self.k4_matrix)245        term_4_t = (d_t4 * a_4_t).sum(dim=-1)246        term_5_t = batch_trace(self.k4_matrix)247        term_6_t = batch_trace(self.k3_matrix)248        term_7_t = batch_diagonal(self.k2_matrix).pow(3).sum(-1)249        term8_t = torch.sum(self.k3_matrix, dim=[-2, -1])250        term9_t = batch_diagonal(self.k2_matrix).pow(2).sum(-1)251        term10_t = batch_trace(self.k2_matrix)252 253        c6_t = (term_1_t - 3 * term_2_t + 9 * term3_t - 6 * term_4_t + 6 * term_5_t - 4 * term_6_t + 4 * term_7_t +254                3 * term8_t - 12 * term9_t + 4 * term10_t)255        return None, (c6_t / 12).unsqueeze(-1).float()256 257    def k_cycles(self, adj_matrix, verbose=False):258        self.adj_matrix = adj_matrix259        self.calculate_kpowers()260 261        k3x, k3y = self.k3_cycle()262        assert (k3x >= -0.1).all()263 264        k4x, k4y = self.k4_cycle()265        assert (k4x >= -0.1).all()266 267        k5x, k5y = self.k5_cycle()268        assert (k5x >= -0.1).all(), k5x269 270        _, k6y = self.k6_cycle()271        assert (k6y >= -0.1).all()272 273        kcyclesx = torch.cat([k3x, k4x, k5x], dim=-1)274        kcyclesy = torch.cat([k3y, k4y, k5y, k6y], dim=-1)275        return kcyclesx, kcyclesy