train_ratio = 0.5
val_ratio = 0.25
test_ratio = 0.25

import numpy as np
from data import load_dataset
import os.path as osp

dataset = 'minesweeper' 
data_dir = osp.join(osp.expanduser('~'), 'datasets/')
data = load_dataset(data_dir, dataname=dataset)
runs = 10

N = data.graph['num_nodes']

train_num = int(N * train_ratio)
valid_num = int(N * val_ratio)

split = []
for r in range(runs):
    perm = np.random.permutation(N)
    
    idx = {}
    idx['train'] = perm[:train_num]
    idx['valid'] = perm[train_num:train_num + valid_num]
    idx['test'] = perm[train_num + valid_num:]

    split.append(idx)

np.save(f"splits/{dataset}-split.npy", split)