gnn/main.go

package main

import (
	"fmt"
	"os"
)

var labelName = []string{
	"Case_Based",
	"Genetic_Algorithms",
	"Neural_Networks",
	"Probabilistic_Methods",
	"Reinforcement_Learning",
	"Rule_Learning",
	"Theory",
}

var nodeLabel = make(map[int]int)

func main() {
	file, err := os.Open("cora/cora.content")
	if err != nil {
		panic(err)
	}
	G := MakeGraph()
	for i := 0; i < 2708; i++ {
		var u int
		fmt.Fscan(file, &u)
		V := MakeVector(1433)
		for j := 0; j < 1433; j++ {
			fmt.Fscan(file, &V[j])
		}
		G.AddNode(u, V)
		var label string
		fmt.Fscan(file, &label)
		for j := 0; j < 7; j++ {
			if labelName[j] == label {
				nodeLabel[u] = j
				break
			}
		}
		if i%100 == 0 {
			fmt.Println("reading node", i)
		}
	}
	file.Close()
	file, err = os.Open("cora/cora.cites")
	if err != nil {
		panic(err)
	}
	for i := 0; i < 5429; i++ {
		var u, v int
		fmt.Fscan(file, &u, &v)
		G.AddEdge(u, v)
	}
	file.Close()
	l := Train(G)
	cnt, tot := 0, 0
	for u := range G.X {
		GetEmbedding(G, u, 2, l)
		id, max := 0, 0.
		for i := 0; i < 7; i++ {
			if l[2].E[0][i] > max {
				id, max = i, l[2].E[0][i]
			}
		}
		if nodeLabel[u] == id {
			cnt++
		}
		tot++
	}
	fmt.Println(cnt, "/", tot)
}