tools/rw-heatmaps: reimplement in go

Signed-off-by: Ivan Valdes <ivan@vald.es>

tools/rw-heatmaps/main.go

@@ -0,0 +1,372 @@
+package main
+
+import (
+	"encoding/csv"
+	"flag"
+	"fmt"
+	"image/color"
+	"io"
+	"log"
+	"math"
+	"os"
+	"sort"
+	"strconv"
+	"strings"
+
+	"gonum.org/v1/plot"
+	"gonum.org/v1/plot/palette"
+	"gonum.org/v1/plot/palette/brewer"
+	"gonum.org/v1/plot/plotter"
+	"gonum.org/v1/plot/vg"
+	"gonum.org/v1/plot/vg/draw"
+	"gonum.org/v1/plot/vg/vgimg"
+)
+
+var (
+	inputFileA   string
+	inputFileB   string
+	title        string
+	zeroCentered bool
+	outputImage  string
+	outputFormat string
+)
+
+func init() {
+	log.SetFlags(0)
+	log.SetPrefix("[INFO] ")
+
+	flag.StringVar(&inputFileA, "input_file_a", "", "first input data file in csv format. (required)")
+	flag.StringVar(&inputFileB, "input_file_b", "", "second input data file in csv format. (optional)")
+	flag.StringVar(&title, "title", "", "plot graph title string (required)")
+	flag.BoolVar(&zeroCentered, "zero-centered", true, "plot the improvement graph with white color represents 0.0")
+	flag.StringVar(&outputImage, "output-image-file", "", "output image filename (required)")
+	flag.StringVar(&outputFormat, "output-format", "png", "output image file format. default: png")
+
+	flag.Parse()
+
+	if inputFileA == "" || title == "" || outputImage == "" {
+		fmt.Println("Missing required arguments.")
+		flag.Usage()
+		os.Exit(2)
+	}
+}
+
+type DataSet struct {
+	Records map[float64][]DataRecord
+	Param   string
+}
+
+type DataRecord struct {
+	ConnSize  int64
+	ValueSize int64
+	Read      float64
+	Write     float64
+}
+
+func loadCSVData(inputFile string) (*DataSet, error) {
+	file, err := os.Open(inputFile)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	reader := csv.NewReader(file)
+	lines, err := reader.ReadAll()
+	if err != nil {
+		return nil, err
+	}
+	dataset := &DataSet{Records: make(map[float64][]DataRecord)}
+	records := dataset.Records
+
+	iters := 0
+	for _, header := range lines[0][4:] {
+		if strings.HasPrefix(header, "iter") {
+			iters++
+		}
+	}
+
+	for _, line := range lines[2:] { // Skip header line.
+		ratio, _ := strconv.ParseFloat(line[1], 64)
+		if _, ok := records[ratio]; !ok {
+			records[ratio] = make([]DataRecord, 0)
+		}
+		connSize, _ := strconv.ParseInt(line[2], 10, 64)
+		valueSize, _ := strconv.ParseInt(line[3], 10, 64)
+
+		readSum := float64(0)
+		writeSum := float64(0)
+
+		for _, v := range line[4 : 4+iters] {
+			splitted := strings.Split(v, ":")
+			readValue, _ := strconv.ParseFloat(splitted[0], 64)
+			readSum += readValue
+
+			writeValue, _ := strconv.ParseFloat(splitted[1], 64)
+			writeSum += writeValue
+		}
+
+		records[ratio] = append(records[ratio], DataRecord{
+			ConnSize:  connSize,
+			ValueSize: valueSize,
+			Read:      readSum / float64(iters),
+			Write:     writeSum / float64(iters),
+		})
+	}
+	dataset.Param = lines[1][iters+4]
+	return dataset, nil
+}
+
+// HeatMapGrid holds X, Y, Z values for a heatmap.
+type HeatMapGrid struct {
+	x, y []float64
+	z    [][]float64 // The Z values should be arranged in a 2D slice.
+}
+
+// Len implements the plotter.GridXYZ interface.
+func (h *HeatMapGrid) Dims() (int, int) {
+	return len(h.x), len(h.y)
+}
+
+// ValueAt returns the value of a grid cell at (c, r).
+// It implements the plotter.GridXYZ interface.
+func (h *HeatMapGrid) Z(c, r int) float64 {
+	return h.z[r][c]
+}
+
+// X returns the coordinate for the column at index c.
+// It implements the plotter.GridXYZ interface.
+func (h *HeatMapGrid) X(c int) float64 {
+	if c >= len(h.x) {
+		panic("index out of range")
+	}
+	return h.x[c]
+}
+
+// Y returns the coordinate for the row at index r.
+// It implements the plotter.GridXYZ interface.
+func (h *HeatMapGrid) Y(r int) float64 {
+	if r >= len(h.y) {
+		panic("index out of range")
+	}
+	return h.y[r]
+}
+
+func uniqueSortedFloats(input []float64) []float64 {
+	unique := make([]float64, 0)
+	seen := make(map[float64]bool)
+
+	for _, value := range input {
+		if !seen[value] {
+			seen[value] = true
+			unique = append(unique, value)
+		}
+	}
+
+	sort.Float64s(unique)
+	return unique
+}
+
+func populateGridAxes(records []DataRecord) ([]float64, []float64) {
+	var xslice, yslice []float64
+
+	for _, record := range records {
+		xslice = append(xslice, float64(record.ConnSize))
+		yslice = append(yslice, float64(record.ValueSize))
+	}
+
+	// Sort and deduplicate the slices
+	xUnique := uniqueSortedFloats(xslice)
+	yUnique := uniqueSortedFloats(yslice)
+
+	return xUnique, yUnique
+}
+
+func plotHeatMaps(title, plotType string, dataset *DataSet) {
+	const rows, cols = 4, 2
+	plots := make([][]*plot.Plot, rows)
+	legends := make([][]plot.Legend, rows)
+	for i := range plots {
+		plots[i] = make([]*plot.Plot, cols)
+		legends[i] = make([]plot.Legend, cols)
+	}
+
+	row, col := 0, 0
+	ratios := make([]float64, 0)
+	for ratio := range dataset.Records {
+		ratios = append(ratios, ratio)
+	}
+	sort.Float64s(ratios)
+	for _, ratio := range ratios {
+		records := dataset.Records[ratio]
+		p, l := plotHeatMap(fmt.Sprintf("R/W Ratio %0.04f", ratio), plotType, records)
+		plots[row][col] = p
+		legends[row][col] = l
+
+		if col++; col == cols {
+			col = 0
+			row++
+		}
+	}
+
+	const width, height = 12 * vg.Inch, 16 * vg.Inch
+	img := vgimg.New(width, height)
+	dc := draw.New(img)
+
+	t := draw.Tiles{
+		Rows:      rows,
+		Cols:      cols,
+		PadX:      vg.Inch * 0.5,
+		PadY:      vg.Inch * 0.5,
+		PadTop:    vg.Inch * 0.5,
+		PadBottom: vg.Inch * 0.5,
+		PadLeft:   vg.Inch * 0.25,
+		PadRight:  vg.Inch * 0.25,
+	}
+
+	canvases := plot.Align(plots, t, dc)
+	for i := 0; i < rows; i++ {
+		for j := 0; j < cols; j++ {
+			if plots[i][j] != nil {
+				l := legends[i][j]
+				r := l.Rectangle(canvases[i][j])
+				legendWidth := r.Max.X - r.Min.X
+				l.YOffs = -plots[i][j].Title.TextStyle.FontExtents().Height // Adjust the legend down a little.
+				l.Draw(canvases[i][j])
+
+				c := draw.Crop(canvases[i][j], 0, -legendWidth-vg.Millimeter, 0, 0) // Make space for the legend.
+				plots[i][j].Draw(c)
+			}
+		}
+	}
+
+	l := plot.NewLegend()
+	l.Add(title)
+	l.Add(dataset.Param)
+	l.Top = true
+	l.Left = true
+	l.Draw(dc)
+
+	fh, err := os.Create(fmt.Sprintf("%s_%s_heatmap.%s", outputImage, plotType, outputFormat))
+	if err != nil {
+		panic(err)
+	}
+	defer fh.Close()
+
+	var w io.WriterTo
+	switch outputFormat {
+	case "png":
+		w = vgimg.PngCanvas{Canvas: img}
+	case "jpeg", "jpg":
+		w = vgimg.PngCanvas{Canvas: img}
+	case "tiff":
+		w = vgimg.TiffCanvas{Canvas: img}
+	}
+	if _, err := w.WriteTo(fh); err != nil {
+		panic(err)
+	}
+}
+
+type pow2Ticks struct{}
+
+func (pow2Ticks) Ticks(min, max float64) []plot.Tick {
+	var t []plot.Tick
+	for i := math.Log2(min); math.Pow(2, i) <= max; i++ {
+		t = append(t, plot.Tick{
+			Value: math.Pow(2, i),
+			Label: fmt.Sprintf("2^%d", int(i)),
+		})
+	}
+	return t
+}
+
+func plotHeatMap(title, plotType string, records []DataRecord) (*plot.Plot, plot.Legend) {
+	p := plot.New()
+	p.X.Scale = plot.LogScale{}
+	p.X.Tick.Marker = pow2Ticks{}
+	p.X.Label.Text = "Connections Amount"
+	p.Y.Scale = plot.LogScale{}
+	p.Y.Tick.Marker = pow2Ticks{}
+	p.Y.Label.Text = "Value Size"
+
+	// Populate X and Y axis data from records
+	xCoords, yCoords := populateGridAxes(records)
+
+	gridData := &HeatMapGrid{
+		x: xCoords,
+		y: yCoords,
+		z: make([][]float64, len(yCoords)),
+	}
+
+	for i := range gridData.z {
+		gridData.z[i] = make([]float64, len(xCoords))
+
+		for j := range gridData.z[i] {
+			recordIndex := i*len(gridData.x) + j
+			if recordIndex >= len(records) {
+				break
+			}
+			record := records[recordIndex]
+			if plotType == "read" {
+				gridData.z[i][j] = record.Read
+			} else {
+				gridData.z[i][j] = record.Write
+			}
+		}
+	}
+
+	colors, _ := brewer.GetPalette(brewer.TypeAny, "YlGnBu", 9)
+	pal := invertedPalette{colors}
+	h := plotter.NewHeatMap(gridData, pal)
+
+	p.Title.Text = title + fmt.Sprintf(" [%.2f, %.2f]", h.Min, h.Max)
+	p.Add(h)
+
+	// Create a legend with the scale.
+	l := plot.NewLegend()
+	thumbs := plotter.PaletteThumbnailers(pal)
+	step := (h.Max - h.Min) / float64(len(thumbs)-1)
+	for i := len(thumbs) - 1; i >= 0; i-- {
+		t := thumbs[i]
+		l.Add(fmt.Sprintf("%.0f", h.Min+step*float64(i)), t)
+	}
+	l.Top = true
+
+	return p, l
+}
+
+// invertedPalette takes an existing palette and inverts it.
+type invertedPalette struct {
+	Base palette.Palette
+}
+
+// Colors returns the sequence of colors in reverse order from the base palette.
+func (p invertedPalette) Colors() []color.Color {
+	baseColors := p.Base.Colors()
+	invertedColors := make([]color.Color, len(baseColors))
+	for i, c := range baseColors {
+		invertedColors[len(baseColors)-i-1] = c
+	}
+	return invertedColors
+}
+
+func main() {
+	//var aRecords []DataRecord
+	//var bRecords []DataRecord
+	//var err error
+
+	aRecords, err := loadCSVData(inputFileA)
+	if err != nil {
+		log.Fatalf("failed to load data from %s: %v\n", inputFileA, err)
+	}
+
+	//	if inputFileB != "" {
+	//		bRecords, err = loadCSVData(inputFileB)
+	//		if err != nil {
+	//			log.Fatalf("failed to load data from %s: %v\n", inputFileB, err)
+	//		}
+	//	}
+
+	//plotHeatMap(title+" Read Plot", "read", aRecords, maxRead)
+	plotHeatMaps(fmt.Sprintf("%s [READ]", title), "read", aRecords)
+	plotHeatMaps(fmt.Sprintf("%s [WRITE]", title), "write", aRecords)
+}