dataset_image.py

# copied from https://github.com/ivam-he/BernNet
# load the image dataset from the `"BernNet: Learning Arbitrary Graph Spectral Filters via Bernstein Approximation" paper
from torch_geometric.data import InMemoryDataset
import torch
from torch_geometric.data.data import Data
import scipy.io as sio
import numpy as np
import matplotlib.pyplot as plt
from torch_geometric.utils import to_scipy_sparse_matrix
import os
from numpy.linalg import eigh
import math

filter_type = ["low", "high", "band", "rejection", "comb", "low_band"]


class TwoDGrid(InMemoryDataset):
    def __init__(self, root="./data/2Dgrid", transform=None, pre_transform=None):
        super(TwoDGrid, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["2Dgrid.mat"]

    @property
    def processed_file_names(self):
        return "data.pt"

    def download(self):
        pass

    def process(self):
        # Read data into huge `Data` list.
        # b = self.processed_paths[0]
        a = sio.loadmat(self.raw_paths[0])  # 'subgraphcount/randomgraph.mat')
        # list of adjacency matrix
        A = a["A"]
        # list of output
        F = a["F"]
        F = F.astype(np.float32)
        # Y=a['Y']
        # Y=Y.astype(np.float32)
        M = a["mask"]
        M = M.astype(np.float32)

        data_list = []
        E = np.where(A > 0)
        edge_index = torch.Tensor(np.vstack((E[0], E[1]))).type(torch.int64)
        x = torch.tensor(F)
        # y=torch.tensor(Y)
        m = torch.tensor(M)

        x_tmp = x[:, 0:1]
        data_list.append(Data(edge_index=edge_index, x=x, x_tmp=x_tmp, m=m))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])


def visualize(y):
    # y=tensor.detach().cpu().numpy()
    y = np.reshape(y, (100, 100))
    plt.imshow(y.T)
    plt.colorbar()
    plt.show()


def myeign(L):
    if os.path.exists("./data/eigenvalues.npy") and os.path.exists(
        "./data/eigenvectors.npy"
    ):
        eigenvalues = np.load("./data/eigenvalues.npy")
        eigenvectors = np.load("./data/eigenvectors.npy")
    else:
        eigenvalues, eigenvectors = eigh(L)
        np.save("./data/eigenvalues.npy", eigenvalues)
        np.save("./data/eigenvectors.npy", eigenvectors)
    return eigenvalues, eigenvectors


def filtering(filter_type, dataset):
    data = dataset[0]
    x = data.x.numpy()

    # print(data.edge_index)
    adj = to_scipy_sparse_matrix(data.edge_index).todense()
    nnodes = adj.shape[0]
    D_vec = np.sum(adj, axis=1).A1
    # print(D_vec.tolist())
    D_vec_invsqrt_corr = 1 / np.sqrt(D_vec)
    D_invsqrt_corr = np.diag(D_vec_invsqrt_corr)
    # print(D_invsqrt_corr)
    L = np.eye(nnodes) - D_invsqrt_corr @ adj @ D_invsqrt_corr
    # print(L)
    eigenvalues, eigenvectors = myeign(L)
    # print(eigenvalues[3])

    # low-pass
    if filter_type == "low":
        value_tmp = [math.exp(-10 * (xxx - 0) ** 2) for xxx in eigenvalues]

    # high-pass
    elif filter_type == "high":
        value_tmp = [1 - math.exp(-10 * (xxx - 0) ** 2) for xxx in eigenvalues]

    # band-pass
    elif filter_type == "band":
        value_tmp = [math.exp(-10 * (xxx - 1) ** 2) for xxx in eigenvalues]

    # band_rejection
    elif filter_type == "rejection":
        value_tmp = [1 - math.exp(-10 * (xxx - 1) ** 2) for xxx in eigenvalues]

    # comb
    elif filter_type == "comb":
        value_tmp = [abs(np.sin(xxx * math.pi)) for xxx in eigenvalues]

    # low_band
    elif filter_type == "low_band":
        y = []
        for i in eigenvalues:
            if i < 0.5:
                y.append(1)
            elif i < 1 and i >= 0.5:
                y.append(math.exp(-100 * (i - 0.5) ** 2))
            else:
                y.append(math.exp(-50 * (i - 1.5) ** 2))
        value_tmp = y

    value_tmp = np.array(value_tmp)
    value_tmp = np.diag(value_tmp)
    # print(value_tmp[5000][5000])

    y = eigenvectors @ value_tmp @ eigenvectors.T @ x
    np.save("y_" + filter_type + ".npy", y)
    return y


def load_img(name):
    ds = TwoDGrid(root="data/2Dgrid", pre_transform=None)
    y = filtering(name, ds)
    y = torch.Tensor(y)
    data = ds[0]
    x = data.x
    ei = data.edge_index
    ea = torch.ones((ei[1].shape[0]))
    mask = data.m.to(torch.long)
    return x, y, ei, ea, mask