scattering_method.py

import torch
import torch.nn as nn
from torch_geometric.utils import add_self_loops, degree


def lazy_random_walk(edge_index, num_nodes):
    row, col = edge_index
    deg = degree(row, num_nodes=num_nodes, dtype=torch.float)
    deg_inv = deg.pow(-1.0)
    deg_inv[deg_inv == float("inf")] = 0
    edge_weight = deg_inv[row]
    return edge_index, edge_weight


def scatter_transform_diffusion(
    x,
    edge_index,
    num_nodes,
    depth=1,
    K=2,
    mode="ADJ",
    th=0.5,
    matrix_scale=0.5,
    matrix_shift=0,
    pruning=False,
):
    """
    Scatter Transform Diffusion Operator

    depth : depth of Scattering Transform
    K     : number of scales per depth (powers of operator)
    mode  : 'ADJ' -> ST, 'LAP', or 'JOINT' -> CST
    """
    print(f"Mode: {mode}, depth={depth}, K={K}, Pruning={pruning}")

    # Build normalized adjacency once
    edge_index, norm = lazy_random_walk(edge_index, num_nodes)
    row, col = edge_index
    adj = torch.sparse_coo_tensor(
        torch.stack([row, col]), norm.view(-1), (num_nodes, num_nodes)
    ).coalesce()

    # Identity
    I = torch.sparse_coo_tensor(
        torch.arange(num_nodes, device=x.device).repeat(2, 1),
        torch.ones(num_nodes, device=x.device),
        (num_nodes, num_nodes),
    )

    # Build both operators in case of JOINT
    if mode == "JOINT":
        K = K // 2
    print(f"K is {K}")
    A_powers, L_powers = {}, {}
    if mode in ["ADJ", "JOINT"]:
        A = matrix_scale * (I + adj)
        A_powers[1] = A
        max_power = 2**K
        p = 1
        while 2**p <= max_power:
            A_powers[2**p] = torch.sparse.mm(
                A_powers[2 ** (p - 1)], A_powers[2 ** (p - 1)]
            )
            p += 1

    if mode in ["LAP", "JOINT"]:
        L = matrix_scale * (I - adj)
        L_powers[1] = L
        max_power = 2**K
        p = 1
        while 2**p <= max_power:
            L_powers[2**p] = torch.sparse.mm(
                L_powers[2 ** (p - 1)], L_powers[2 ** (p - 1)]
            )
            p += 1

    # Helper to compute features for one operator
    def compute_branch(prev_x, powers, operator_type):
        branch_features = []
        for j in range(1, depth + 1):
            print(f"[{operator_type}] Depth {j}, prev_x: {prev_x.shape}")
            prev_energy = torch.norm(prev_x)
            scale_sig = []
            futures = []
            for k in range(K):
                if operator_type == "ADJ":
                    futures.append(
                        torch.jit.fork(
                            lambda k: (
                                torch.abs(
                                    torch.sparse.mm(powers[2 ** (k - 1)], prev_x)
                                    - torch.sparse.mm(powers[2**k], prev_x)
                                )
                                if k > 0
                                else torch.abs(
                                    prev_x - torch.sparse.mm(powers[1], prev_x)
                                )
                            ),
                            k,
                        )
                    )
                    # futures.append(torch.jit.fork(
                    #     lambda k: torch.abs(torch.sparse.mm(powers[2**(k-1)], prev_x) -
                    #                torch.sparse.mm(powers[2**k], prev_x)),
                    #     k
                    # ))
                else:  # LAP
                    futures.append(
                        torch.jit.fork(
                            lambda k: (
                                torch.abs(
                                    torch.sparse.mm(powers[2 ** (k - 1)], prev_x)
                                    - torch.sparse.mm(powers[2**k], prev_x)
                                )
                                if k > 0
                                else torch.abs(
                                    prev_x - torch.sparse.mm(powers[1], prev_x)
                                )
                            ),
                            k,
                        )
                    )
                    # futures.append(torch.jit.fork(
                    #     lambda k: torch.abs(torch.sparse.mm(powers[2**(k-1)], prev_x) -
                    #                         torch.sparse.mm(powers[2**k], prev_x)),
                    #     k
                    # ))
            scale_sig_fut = [torch.jit.wait(f) for f in futures]
            for out in scale_sig_fut:
                if pruning:
                    out_energy = torch.norm(out)
                    if out_energy / prev_energy > th:
                        scale_sig.append(out)
                else:
                    scale_sig.append(out)
            scale_feature = torch.hstack(scale_sig)
            branch_features.append(scale_feature)
            prev_x = scale_feature  # continue chain
        return branch_features

    # Run depending on mode
    features = [x]

    if mode == "ADJ":
        features += compute_branch(x, A_powers, "ADJ")

    elif mode == "LAP":
        features += compute_branch(x, L_powers, "LAP")

    elif mode == "JOINT":
        adj_feats = compute_branch(x, A_powers, "ADJ")
        lap_feats = compute_branch(x, L_powers, "LAP")
        features += adj_feats + lap_feats

    return torch.cat(features, dim=1)


class LogReg(nn.Module):
    def __init__(self, dim, n_class):
        super(LogReg, self).__init__()
        self.fc = nn.Linear(dim, n_class)

    def forward(self, x):
        x = self.fc(x)
        return x