Confusion matrix!

Author: ammario

go.mod

@@ -174,6 +174,7 @@ require (
 	github.com/coder/flog v1.0.0 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect
 	github.com/gocarina/gocsv v0.0.0-20230123225133-763e25b40669 // indirect
+	github.com/goml/gobrain v0.0.0-20201212123421-2e2d98ca8249 // indirect
 	github.com/hashicorp/go-cty v1.4.1-0.20200414143053-d3edf31b6320 // indirect
 	github.com/hashicorp/go-hclog v1.2.1 // indirect
 	github.com/hashicorp/go-plugin v1.4.4 // indirect

go.sum

@@ -860,6 +860,8 @@ github.com/golangci/maligned v0.0.0-20180506175553-b1d89398deca/go.mod h1:tvlJhZ
 github.com/golangci/misspell v0.3.5/go.mod h1:dEbvlSfYbMQDtrpRMQU675gSDLDNa8sCPPChZ7PhiVA=
 github.com/golangci/revgrep v0.0.0-20210930125155-c22e5001d4f2/go.mod h1:LK+zW4MpyytAWQRz0M4xnzEk50lSvqDQKfx304apFkY=
 github.com/golangci/unconvert v0.0.0-20180507085042-28b1c447d1f4/go.mod h1:Izgrg8RkN3rCIMLGE9CyYmU9pY2Jer6DgANEnZ/L/cQ=
+github.com/goml/gobrain v0.0.0-20201212123421-2e2d98ca8249 h1:Xst86cFqcrNSUxId+2A/o3NZUwm7H+H6bUcehmD2t4o=
+github.com/goml/gobrain v0.0.0-20201212123421-2e2d98ca8249/go.mod h1:imJK2QRE3080lm54nE96nRHXv7EvvVMCbF5N2QdGa9I=
 github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/btree v1.0.1 h1:gK4Kx5IaGY9CD5sPJ36FHiBJ6ZXl0kilRiiCj+jdYp4=

scripts/aistart/loadcsv.go

@@ -27,16 +27,23 @@ type trainingRow struct {
 	// HourOfDay ranges from 0 to 23
 	HourOfDay int `csv:"hour"`
 	// Day of Week ranges from 0 to 6
-	DayOfWeek int  `csv:"day"`
-	Used      bool `csv:"used"`
+	DayOfWeek int `csv:"day"`
+	Used      int `csv:"used"`
+}
+
+func (t trainingRow) vectorize() vector {
+	return [][]float64{
+		{float64(t.HourOfDay) / 23, float64(t.DayOfWeek / 6)},
+		{float64(t.Used)},
+	}
 }
 
 type dbRow struct {
 	Time        time.Time
 	WorkspaceID string
 }
 
-func (db dbRow) convert(used bool) trainingRow {
+func (db dbRow) convert(used int) trainingRow {
 	return trainingRow{
 		WorkspaceID: db.WorkspaceID,
 		HourOfDay:   db.Time.Hour(),
@@ -66,12 +73,12 @@ func generateTrainingRows(rs []dbRow) []trainingRow {
 						WorkspaceID: wid,
 						HourOfDay:   last.Hour(),
 						DayOfWeek:   int(last.Weekday()),
-						Used:        false,
+						Used:        0,
 					})
 				}
 			}
 		}
-		trainingRows = append(trainingRows, r.convert(true))
+		trainingRows = append(trainingRows, r.convert(1))
 	}
 
 	return trainingRows
@@ -124,8 +131,8 @@ func loadTrainingCSV() *cobra.Command {
 			flog.Info("loaded %v rows", len(rs))
 			trainingRows := generateTrainingRows(rs)
 			flog.Info("generated %v training rows", len(trainingRows))
-			gocsv.Marshal(trainingRows, os.Stdout)
-			return nil
+			err = gocsv.Marshal(trainingRows, os.Stdout)
+			return err
 		},
 	}
 }

scripts/aistart/main.go

@@ -16,6 +16,7 @@ func main() {
 	}
 
 	cmd.AddCommand(loadTrainingCSV())
+	cmd.AddCommand(train())
 
 	cmd, err := cmd.ExecuteC()
 	if err != nil {

scripts/aistart/train.go

@@ -0,0 +1,89 @@
+package main
+
+import (
+	"fmt"
+	"math"
+	"math/rand"
+	"os"
+	"text/tabwriter"
+
+	"github.com/gocarina/gocsv"
+	"github.com/goml/gobrain"
+	"github.com/spf13/cobra"
+
+	"github.com/coder/flog"
+)
+
+type vector [][]float64
+
+type pattern []vector
+
+func (p pattern) floats() [][][]float64 {
+	var r [][][]float64
+	for _, v := range p {
+		r = append(r, [][]float64(v))
+	}
+	return r
+}
+
+func vectorizeTrainingRows(rs []trainingRow) pattern {
+	var p pattern
+	for _, r := range rs {
+		p = append(p, r.vectorize())
+	}
+	return p
+}
+
+func splitTrainTest(rat float64, p pattern) (train, test pattern) {
+	perms := rand.Perm(len(p))
+	for i, v := range p {
+		if float64(perms[i])/float64(len(p)) > rat {
+			test = append(test, v)
+		} else {
+			train = append(train, v)
+		}
+	}
+	return train, test
+}
+
+func train() *cobra.Command {
+	return &cobra.Command{
+		Use: "train",
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			var rs []trainingRow
+
+			err := gocsv.Unmarshal(os.Stdin, &rs)
+			if err != nil {
+				return err
+			}
+
+			all := vectorizeTrainingRows(rs)
+
+			train, test := splitTrainTest(0.5, all)
+
+			flog.Info("split train test: %v/%v", len(train), len(test))
+
+			ff := &gobrain.FeedForward{}
+			ff.Init(2, 2, 1)
+			ff.Train(train.floats(), 50, 0.001, 0.4, true)
+			var (
+				// confusionMatrix has actual values in the first index with
+				// predicted values in the second.
+				confusionMatrix [2][2]int
+			)
+			for _, v := range train {
+				want := v[1][0]
+				gotArr := ff.Update(v[0])
+				got := gotArr[0]
+				confusionMatrix[0][int(math.Round(want))]++
+				confusionMatrix[1][int(math.Round(got))]++
+			}
+			twr := tabwriter.NewWriter(os.Stderr, 0, 4, 3, ' ', 0)
+			_, _ = fmt.Fprintf(twr, "-\tOff\tOn\n")
+			_, _ = fmt.Fprintf(twr, "Actual\t%v\t%v\n", confusionMatrix[0][0], confusionMatrix[0][1])
+			_, _ = fmt.Fprintf(twr, "Predicted\t%v\t%v\n", confusionMatrix[1][0], confusionMatrix[1][1])
+			twr.Flush()
+			return nil
+		},
+	}
+}