Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -42,18 +42,18 @@ Now, you need to preprocess Elliptic2 to obtain the specific input format for ea

#### GLASS and GNNSeg

To preprocess for GLASS and GNNSeg, edit `DATAPATH` in `preprocess_glass.py` and run
To preprocess for GLASS and GNNSeg, run (use `--datapath` if your data is not in `./dataset`)
```
python preprocess_glass.py
```
which will produce three files: `edge_list.txt`, `subgraph.pth`, `n2id.pkl`.

The first two files are the required inputs for GLASS, which should be move into `GLASS/dataset/elliptic`. The file `n2id.pkl` will store a dictionary that converts node id in Ellptic2 to node id in GLASS for reference. Note that the train/val split is fixed given the random seeds, but you can change the split size at the top of the python script.
The first two files are the required inputs for GLASS, which should be move into `GLASS/dataset/elliptic`. The file `n2id.pkl` will store a dictionary that converts node id in Ellptic2 to node id in GLASS for reference. Note that the train/val/test split is fixed (roughly 80/10/10) by the deterministic per-subgraph assignment in the script.


#### Sub2vec

To preprocess for Sub2Vec, edit `DATAPATH` in `preprocess_sub2vec.py` and run
To preprocess for Sub2Vec, run (use `--datapath` if needed)
```
python preprocess_sub2vec.py
```
Expand Down
21 changes: 12 additions & 9 deletions preprocess_glass.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,18 @@
import random
import time
import pickle
import argparse
import os

#Parameters
DATAPATH = "./dataset/"
train = 0.8 #percentage of training subgraph
val = 0.1 #percentage of validation subgraph
parser = argparse.ArgumentParser(description="Preprocess the Elliptic2 dataset for GLASS and GNNSeg.")
parser.add_argument("--datapath", default=os.environ.get("ELLIPTIC2_DATAPATH", "./dataset"),
help="Folder with the unzipped Elliptic2 CSV files (default: ./dataset).")
DATAPATH = parser.parse_args().datapath

#Read in nodes
start = time.time()
feat = pd.read_csv(DATAPATH+"/background_nodes.csv")
feat = pd.read_csv(os.path.join(DATAPATH, "background_nodes.csv"))
print(list(feat.columns[1:]))
print(feat.head())
print("load background node time", time.time()-start)
Expand All @@ -31,7 +34,7 @@
pickle.dump(n2id, fp)

#Read in edge list
edge = pd.read_csv(DATAPATH+"/background_edges.csv",usecols=["clId1","clId2"])
edge = pd.read_csv(os.path.join(DATAPATH, "background_edges.csv"),usecols=["clId1","clId2"])
print("load background edge time", time.time()-start)
print("total number of edge: ",len(edge))
start = time.time()
Expand All @@ -42,7 +45,7 @@
if n2id[c1] > maxid:
print("WARNING NODE OUT OF RANGE:",c1,n2id[c1])
if n2id[c2] > maxid:
print("WARNING NODE OUT OF RANGE:",c2,n2id[c1])
print("WARNING NODE OUT OF RANGE:",c2,n2id[c2])
file.write(str(n2id[c1])+" "+str(n2id[c2])+"\n")
file.close()
print("time to store edgelist", time.time()-start)
Expand All @@ -51,9 +54,9 @@
#Read in Subgraph

start = time.time()
cc = pd.read_csv(DATAPATH+"connected_components.csv")
edge = pd.read_csv(DATAPATH+"edges.csv")
node = pd.read_csv(DATAPATH+"nodes.csv")
cc = pd.read_csv(os.path.join(DATAPATH, "connected_components.csv"))
edge = pd.read_csv(os.path.join(DATAPATH, "edges.csv"))
node = pd.read_csv(os.path.join(DATAPATH, "nodes.csv"))
print("load rest time", time.time()-start)
start = time.time()

Expand Down
12 changes: 8 additions & 4 deletions preprocess_sub2vec.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,14 @@
import torch
import pickle
import numpy as np
import os
import os
import argparse

#Parameters
DATAPATH = "./dataset/"
parser = argparse.ArgumentParser(description="Preprocess the Elliptic2 dataset for Sub2Vec.")
parser.add_argument("--datapath", default=os.environ.get("ELLIPTIC2_DATAPATH", "./dataset"),
help="Folder with the unzipped Elliptic2 CSV files (default: ./dataset).")
DATAPATH = parser.parse_args().datapath

#Read in Node ID
n2id = {}
Expand All @@ -14,13 +18,13 @@

#Read in Subgraph ID
cc2id = {}
cc = pd.read_csv(DATAPATH+"/connected_components.csv")
cc = pd.read_csv(os.path.join(DATAPATH, "connected_components.csv"))
for row in cc.itertuples(index=True):
cc2id[int(row[1])] =int(row[0])

#Read in Nodes in Subgraph
sub = {}
node = pd.read_csv(DATAPATH+"/nodes.csv")
node = pd.read_csv(os.path.join(DATAPATH, "nodes.csv"))
for row in node.itertuples(index=False):
if cc2id[int(row[1])] in sub.keys():
sub[cc2id[int(row[1])]].append(n2id[int(row[0])])
Expand Down