diff --git a/tools/rw-heatmaps/main.go b/tools/rw-heatmaps/main.go new file mode 100644 index 000000000000..ebe751ee1c82 --- /dev/null +++ b/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) +}