Add heatmap plots1.1.8

5 files changed, 210 insertions, 6 deletions

diff --git a/pkg/data/heatmap.go b/pkg/data/heatmap.go
new file mode 100644
index 0000000..7a0aa6f
--- /dev/null
+++ b/pkg/data/heatmap.go
@@ -0,0 +1,68 @@
+package data
+
+import "math"
+
+type Heatmap struct {
+	X, Y       []Bin
+	Z          [][]float64
+	MinX, MaxX uint64
+	MinY, MaxY uint64
+	MinZ, MaxZ uint64
+}
+
+func NewHeatmap(x, y []Bin, z [][]uint64) *Heatmap {
+	h := new(Heatmap)
+	h.X, h.Y = x, y
+	h.Z = make([][]float64, len(z))
+	h.MinX, h.MinY, h.MinZ = math.MaxUint64, math.MaxUint64, math.MaxUint64
+	h.MaxX, h.MaxY, h.MaxZ = 0, 0, 0
+	for _, b := range x {
+		if b.Count > h.MaxX {
+			h.MaxX = b.Count
+		}
+		if b.Count < h.MinX {
+			h.MinX = b.Count
+		}
+	}
+	for _, b := range x {
+		b.CountNorm = float64(b.Count-h.MinX) / float64(h.MaxX-h.MinX)
+	}
+	for _, b := range y {
+		if b.Count > h.MaxY {
+			h.MaxY = b.Count
+		}
+		if b.Count < h.MinY {
+			h.MinY = b.Count
+		}
+	}
+	for _, b := range y {
+		b.CountNorm = float64(b.Count-h.MinY) / float64(h.MaxY-h.MinY)
+	}
+	for i := range z {
+		h.Z[i] = make([]float64, len(z[i]))
+		for _, b := range z[i] {
+			if b > h.MaxZ {
+				h.MaxZ = b
+			}
+			if b < h.MinZ {
+				h.MinZ = b
+			}
+		}
+	}
+	for i := range z {
+		for j := range z[i] {
+			h.Z[i][j] = float64(z[i][j]-h.MinZ) / float64(h.MaxZ-h.MinZ)
+		}
+	}
+	if h.MaxX == 0 {
+		h.MaxX = 1
+	}
+	if h.MaxY == 0 {
+		h.MaxY = 1
+	}
+	if h.MaxZ == 0 {
+		h.MaxZ = 1
+	}
+
+	return h
+}
diff --git a/pkg/data/histogram.go b/pkg/data/histogram.go
index e16d000..be72142 100644
--- a/pkg/data/histogram.go
+++ b/pkg/data/histogram.go
@@ -23,6 +23,7 @@ type Bin struct {
 	Right          float64
 	RightInclusive bool
 	Count          uint64
+	CountNorm      float64
 }
 
 func (b *Bin) String() string {
@@ -38,16 +39,16 @@ type Bins struct {
 	numPoints int
 }
 
-func (b *Bins) ChooseSqrt() {
-	b.Number = int(math.Sqrt(float64(b.numPoints)))
+func BinsSqrt(numPoints int) int {
+	return int(math.Sqrt(float64(numPoints)))
 }
 
-func (b *Bins) ChooseSturges() {
-	b.Number = int(math.Ceil(math.Log2(float64(b.numPoints))) + 1)
+func BinsSturges(numPoints int) int {
+	return int(math.Ceil(math.Log2(float64(numPoints))) + 1)
 }
 
-func (b *Bins) ChooseRice() {
-	b.Number = int(math.Ceil(2 * math.Pow(float64(b.numPoints), 1.0/3.0)))
+func BinsRice(numPoints int) int {
+	return int(math.Ceil(2 * math.Pow(float64(numPoints), 1.0/3.0)))
 }
 
 func NewBins(points []float64) *Bins {
@@ -106,3 +107,42 @@ func Histogram(points []float64, bins *Bins) (out []Bin) {
 	}
 	return
 }
+
+type Bins2D struct {
+	X *Bins
+	Y *Bins
+}
+
+func NewBins2D(points [][2]float64) *Bins2D {
+	bins := new(Bins2D)
+	xs := make([]float64, len(points))
+	ys := make([]float64, len(points))
+	for i := range points {
+		xs[i] = points[i][0]
+		ys[i] = points[i][1]
+	}
+	bins.X = NewBins(xs)
+	bins.Y = NewBins(ys)
+	return bins
+}
+
+func Histogram2D(points [][2]float64, bins *Bins2D) (x, y []Bin, z [][]uint64) {
+	x = bins.X.All()
+	y = bins.Y.All()
+	z = make([][]uint64, len(y))
+	for _, b := range x {
+		b.Count = 0
+	}
+	for i, b := range y {
+		z[i] = make([]uint64, len(x))
+		b.Count = 0
+	}
+	for _, p := range points {
+		i := bins.X.Point(p[0])
+		j := bins.Y.Point(p[1])
+		x[i].Count++
+		y[j].Count++
+		z[i][j]++
+	}
+	return
+}
diff --git a/pkg/draw/heatmap.go b/pkg/draw/heatmap.go
new file mode 100644
index 0000000..7f913fd
--- /dev/null
+++ b/pkg/draw/heatmap.go
@@ -0,0 +1,23 @@
+package draw
+
+type Heatmap struct{ *Buffer }
+
+func (b *Heatmap) Size() Box { return b.Box }
+
+var shades = []rune(" ·░▒▒▒▒▓▓▓▓█")
+var nonZeroShade = 1.0 / float64(len(shades)-1)
+
+func (b *Heatmap) Set(y, x int, fill float64) {
+	if fill < 0.0 {
+		fill = 0.0
+	}
+	if fill > 0.0 && fill < nonZeroShade {
+		fill = nonZeroShade
+	}
+	if fill > 1.0 {
+		fill = 1.0
+	}
+	b.Buffer.Set(y, x, shades[int(fill*float64(len(shades)-1))])
+}
+
+func (b *Heatmap) Clear() { b.Fill(shades[0]) }
diff --git a/pkg/jsonpath/jsonpath.go b/pkg/jsonpath/jsonpath.go
index 7ed9b16..6b11660 100644
--- a/pkg/jsonpath/jsonpath.go
+++ b/pkg/jsonpath/jsonpath.go
@@ -26,6 +26,7 @@ import (
 
 type JSONPath struct {
 	name       string
+	String     string
 	parser     *Parser
 	stack      [][]reflect.Value // push and pop values in different scopes
 	cur        []reflect.Value   // current scope values
@@ -56,6 +57,7 @@ func (j *JSONPath) AllowMissingKeys(allow bool) *JSONPath {
 // Parse parses the given template and returns an error.
 func (j *JSONPath) Parse(text string) error {
 	var err error
+	j.String = text
 	j.parser, err = Parse(j.name, text)
 	return err
 }
diff --git a/pkg/plot/heatmap.go b/pkg/plot/heatmap.go
new file mode 100644
index 0000000..2667b40
--- /dev/null
+++ b/pkg/plot/heatmap.go
@@ -0,0 +1,71 @@
+package plot
+
+import (
+	"bytes"
+
+	"github.com/sgreben/jp/pkg/data"
+	"github.com/sgreben/jp/pkg/draw"
+)
+
+// HeatMap is a heatmap
+type HeatMap struct{ draw.Heatmap }
+
+// NewHeatMap returns a new line chart
+func NewHeatMap(buffer *draw.Buffer) *HeatMap { return &HeatMap{draw.Heatmap{buffer}} }
+
+func (c *HeatMap) drawAxes(paddingX, paddingY int, minX, maxX, minY, maxY float64) {
+	buffer := c.GetBuffer()
+	// X axis
+	buffer.SetRow(1, paddingX, buffer.Width, draw.HorizontalLine)
+	// Y axis
+	buffer.SetColumn(1, buffer.Height, paddingX, draw.VerticalLine)
+	// Corner
+	buffer.Set(1, paddingX, draw.CornerBottomLeft)
+	// Labels
+	buffer.WriteRight(1, 1, Ff(minY))
+	buffer.WriteLeft(buffer.Height-1, paddingX, Ff(maxY))
+	buffer.WriteRight(0, paddingX, Ff(minX))
+	buffer.WriteLeft(0, buffer.Width, Ff(maxX))
+}
+
+// Draw implements Chart
+func (c *HeatMap) Draw(heatmap *data.Heatmap) string {
+	var scaleY, scaleX float64
+
+	minX := heatmap.X[0].LeftInclusive
+	maxX := heatmap.X[len(heatmap.X)-1].Right
+	minY := heatmap.Y[0].LeftInclusive
+	maxY := heatmap.Y[len(heatmap.Y)-1].Right
+	minLabelWidth := len(Ff(minY))
+	maxLabelWidth := len(Ff(maxY))
+
+	paddingX := minLabelWidth + 1
+	paddingY := 2
+	if minLabelWidth < maxLabelWidth {
+		paddingX = maxLabelWidth + 1
+	}
+	chartWidth := c.Size().Width - (paddingX + 1)
+	chartHeight := c.Size().Height - paddingY
+	scaleX = float64(chartWidth) / (maxX - minX)
+	scaleY = float64(chartHeight) / (maxY - minY)
+
+	for i := range heatmap.Z {
+		for j := range heatmap.Z[i] {
+			x0 := int((heatmap.X[j].LeftInclusive-minX)*scaleX + float64(paddingX+1))
+			y0 := int((heatmap.Y[i].LeftInclusive-minY)*scaleY + float64(paddingY))
+			x1 := int((heatmap.X[j].Right-minX)*scaleX + float64(paddingX+1))
+			y1 := int((heatmap.Y[i].Right-minY)*scaleY + float64(paddingY))
+			z := heatmap.Z[i][j]
+			for x := x0; x < x1; x++ {
+				for y := y0; y < y1; y++ {
+					c.Set(y, x, z)
+				}
+			}
+		}
+	}
+	c.drawAxes(paddingX, paddingY, minX, maxX, minY, maxY)
+
+	b := bytes.NewBuffer(nil)
+	c.GetBuffer().Render(b)
+	return b.String()
+}