forked from dustin/go-heatmap
/
heatmap.go
165 lines (138 loc) · 3.78 KB
/
heatmap.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
// Generate heatmaps.
package heatmap
import (
"image"
"image/color"
"image/draw"
"math"
"sync"
)
// A data point to be plotted.
// These are all normalized to use the maximum amount of
// space available in the output image.
type DataPoint interface {
X() float64
Y() float64
}
type apoint struct {
x float64
y float64
}
func (a apoint) X() float64 {
return a.x
}
func (a apoint) Y() float64 {
return a.y
}
// Construct a simple datapoint
func P(x, y float64) DataPoint {
return apoint{x, y}
}
type limits struct {
Min DataPoint
Max DataPoint
}
func (l limits) Dx() float64 {
return l.Max.X() - l.Min.X()
}
func (l limits) Dy() float64 {
return l.Max.Y() - l.Min.Y()
}
// Draw a heatmap.
//
// size is the size of the image to crate
// dotSize is the impact size of each point on the output
// opacity is the alpha value (0-255) of the impact of the image overlay
// scheme is the color palette to choose from the overlay
func Heatmap(size image.Rectangle, points []DataPoint, dotSize int, opacity uint8,
scheme []color.Color) image.Image {
dot := mkDot(float64(dotSize))
limits := findLimits(points)
// Draw black/alpha into the image
bw := image.NewRGBA(size)
placePoints(size, limits, bw, points, dot)
rv := image.NewRGBA(size)
// Then we transplant the pixels one at a time pulling from our color map
warm(rv, bw, opacity, scheme)
return rv
}
func placePoints(size image.Rectangle, limits limits,
bw *image.RGBA, points []DataPoint, dot draw.Image) {
for _, p := range points {
limits.placePoint(p, bw, dot)
}
}
func warm(out, in draw.Image, opacity uint8, colors []color.Color) {
bounds := in.Bounds()
collen := float64(len(colors))
wg := sync.WaitGroup{}
wg.Add(bounds.Dx())
for x := bounds.Min.X; x < bounds.Max.X; x++ {
go func(x int) {
defer wg.Done()
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
col := in.At(x, y)
_, _, _, alpha := col.RGBA()
percent := float64(alpha) / float64(0xffff)
var outcol color.Color
if percent == 0 {
outcol = color.Transparent
} else {
template := colors[int((collen-1)*(1.0-percent))]
tr, tg, tb, ta := template.RGBA()
ta /= 256
outalpha := uint8(float64(ta) *
(float64(opacity) / 256.0))
outcol = color.NRGBA{
uint8(tr / 256),
uint8(tg / 256),
uint8(tb / 256),
uint8(outalpha)}
}
out.Set(x, y, outcol)
}
}(x)
}
wg.Wait()
}
func findLimits(points []DataPoint) limits {
minx, miny := points[0].X(), points[0].Y()
maxx, maxy := minx, miny
for _, p := range points {
minx = math.Min(p.X(), minx)
miny = math.Min(p.Y(), miny)
maxx = math.Max(p.X(), maxx)
maxy = math.Max(p.Y(), maxy)
}
return limits{apoint{minx, miny}, apoint{maxx, maxy}}
}
func mkDot(size float64) draw.Image {
i := image.NewRGBA(image.Rect(0, 0, int(size), int(size)))
md := 0.5 * math.Sqrt(math.Pow(float64(size)/2.0, 2)+math.Pow((float64(size)/2.0), 2))
for x := float64(0); x < size; x++ {
for y := float64(0); y < size; y++ {
d := math.Sqrt(math.Pow(x-size/2.0, 2) + math.Pow(y-size/2.0, 2))
if d < md {
rgbVal := uint8(200.0*d/md + 50.0)
rgba := color.NRGBA{0, 0, 0, 255 - rgbVal}
i.Set(int(x), int(y), rgba)
}
}
}
return i
}
func (l limits) translate(p DataPoint, i draw.Image, dotsize int) (rv image.Point) {
// Normalize to 0-1
x := float64(p.X()-l.Min.X()) / float64(l.Dx())
y := float64(p.Y()-l.Min.Y()) / float64(l.Dy())
// And remap to the image
rv.X = int(x * float64((i.Bounds().Max.X - dotsize)))
rv.Y = int((1.0 - y) * float64((i.Bounds().Max.Y - dotsize)))
return
}
func (l limits) placePoint(p DataPoint, i, dot draw.Image) {
pos := l.translate(p, i, dot.Bounds().Max.X)
dotw, doth := dot.Bounds().Max.X, dot.Bounds().Max.Y
draw.Draw(i, image.Rect(pos.X, pos.Y, pos.X+dotw, pos.Y+doth), dot,
image.ZP, draw.Over)
}