275 lines
6.8 KiB
Go
275 lines
6.8 KiB
Go
// Copyright ©2015 The Gonum Authors. All rights reserved.
|
|
// Use of this source code is governed by a BSD-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
package plotter
|
|
|
|
import (
|
|
"image"
|
|
"image/color"
|
|
"math"
|
|
|
|
"gonum.org/v1/plot"
|
|
"gonum.org/v1/plot/palette"
|
|
"gonum.org/v1/plot/vg"
|
|
"gonum.org/v1/plot/vg/draw"
|
|
)
|
|
|
|
// GridXYZ describes three dimensional data where the X and Y
|
|
// coordinates are arranged on a rectangular grid.
|
|
type GridXYZ interface {
|
|
// Dims returns the dimensions of the grid.
|
|
Dims() (c, r int)
|
|
|
|
// Z returns the value of a grid value at (c, r).
|
|
// It will panic if c or r are out of bounds for the grid.
|
|
Z(c, r int) float64
|
|
|
|
// X returns the coordinate for the column at the index c.
|
|
// It will panic if c is out of bounds for the grid.
|
|
X(c int) float64
|
|
|
|
// Y returns the coordinate for the row at the index r.
|
|
// It will panic if r is out of bounds for the grid.
|
|
Y(r int) float64
|
|
}
|
|
|
|
// HeatMap implements the Plotter interface, drawing
|
|
// a heat map of the values in the GridXYZ field.
|
|
type HeatMap struct {
|
|
GridXYZ GridXYZ
|
|
|
|
// Palette is the color palette used to render
|
|
// the heat map. Palette must not be nil or
|
|
// return a zero length []color.Color.
|
|
Palette palette.Palette
|
|
|
|
// Underflow and Overflow are colors used to fill
|
|
// heat map elements outside the dynamic range
|
|
// defined by Min and Max.
|
|
Underflow color.Color
|
|
Overflow color.Color
|
|
|
|
// NaN is the color used to fill heat map elements
|
|
// that are NaN or do not map to a unique palette
|
|
// color.
|
|
NaN color.Color
|
|
|
|
// Min and Max define the dynamic range of the
|
|
// heat map.
|
|
Min, Max float64
|
|
|
|
// Rasterized indicates whether the heatmap
|
|
// should be produced using raster-based drawing.
|
|
Rasterized bool
|
|
}
|
|
|
|
// NewHeatMap creates as new heat map plotter for the given data,
|
|
// using the provided palette. If g has Min and Max methods that return
|
|
// a float, those returned values are used to set the respective HeatMap
|
|
// fields. If the returned HeatMap is used when Min is greater than Max,
|
|
// the Plot method will panic.
|
|
func NewHeatMap(g GridXYZ, p palette.Palette) *HeatMap {
|
|
var min, max float64
|
|
type minMaxer interface {
|
|
Min() float64
|
|
Max() float64
|
|
}
|
|
switch g := g.(type) {
|
|
case minMaxer:
|
|
min, max = g.Min(), g.Max()
|
|
default:
|
|
min, max = math.Inf(1), math.Inf(-1)
|
|
c, r := g.Dims()
|
|
for i := 0; i < c; i++ {
|
|
for j := 0; j < r; j++ {
|
|
v := g.Z(i, j)
|
|
if math.IsNaN(v) {
|
|
continue
|
|
}
|
|
min = math.Min(min, v)
|
|
max = math.Max(max, v)
|
|
}
|
|
}
|
|
}
|
|
|
|
return &HeatMap{
|
|
GridXYZ: g,
|
|
Palette: p,
|
|
Min: min,
|
|
Max: max,
|
|
}
|
|
}
|
|
|
|
// Plot implements the Plot method of the plot.Plotter interface.
|
|
func (h *HeatMap) Plot(c draw.Canvas, plt *plot.Plot) {
|
|
if h.Rasterized {
|
|
h.plotRasterized(c, plt)
|
|
} else {
|
|
h.plotVectorized(c, plt)
|
|
}
|
|
}
|
|
|
|
// plotRasterized plots the heatmap using raster-based drawing.
|
|
func (h *HeatMap) plotRasterized(c draw.Canvas, plt *plot.Plot) {
|
|
cols, rows := h.GridXYZ.Dims()
|
|
img := image.NewRGBA64(image.Rectangle{
|
|
Min: image.Point{X: 0, Y: 0},
|
|
Max: image.Point{X: cols, Y: rows},
|
|
})
|
|
|
|
pal := h.Palette.Colors()
|
|
ps := float64(len(pal)-1) / (h.Max - h.Min)
|
|
for i := 0; i < cols; i++ {
|
|
for j := 0; j < rows; j++ {
|
|
var col color.Color
|
|
switch v := h.GridXYZ.Z(i, j); {
|
|
case v < h.Min:
|
|
col = h.Underflow
|
|
case v > h.Max:
|
|
col = h.Overflow
|
|
case math.IsNaN(v), math.IsInf(ps, 0):
|
|
col = h.NaN
|
|
default:
|
|
col = pal[int((v-h.Min)*ps+0.5)] // Apply palette scaling.
|
|
}
|
|
|
|
if col != nil {
|
|
img.Set(i, rows-j-1, col)
|
|
}
|
|
}
|
|
}
|
|
|
|
xmin, xmax, ymin, ymax := h.DataRange()
|
|
pImg := NewImage(img, xmin, ymin, xmax, ymax)
|
|
pImg.Plot(c, plt)
|
|
}
|
|
|
|
// plotVectorized plots the heatmap using vector-based drawing.
|
|
func (h *HeatMap) plotVectorized(c draw.Canvas, plt *plot.Plot) {
|
|
if h.Min > h.Max {
|
|
panic("contour: invalid Z range: min greater than max")
|
|
}
|
|
pal := h.Palette.Colors()
|
|
if len(pal) == 0 {
|
|
panic("heatmap: empty palette")
|
|
}
|
|
// ps scales the palette uniformly across the data range.
|
|
ps := float64(len(pal)-1) / (h.Max - h.Min)
|
|
|
|
trX, trY := plt.Transforms(&c)
|
|
|
|
var pa vg.Path
|
|
cols, rows := h.GridXYZ.Dims()
|
|
for i := 0; i < cols; i++ {
|
|
var right, left float64
|
|
switch i {
|
|
case 0:
|
|
if cols == 1 {
|
|
right = 0.5
|
|
} else {
|
|
right = (h.GridXYZ.X(1) - h.GridXYZ.X(0)) / 2
|
|
}
|
|
left = -right
|
|
case cols - 1:
|
|
right = (h.GridXYZ.X(cols-1) - h.GridXYZ.X(cols-2)) / 2
|
|
left = -right
|
|
default:
|
|
right = (h.GridXYZ.X(i+1) - h.GridXYZ.X(i)) / 2
|
|
left = -(h.GridXYZ.X(i) - h.GridXYZ.X(i-1)) / 2
|
|
}
|
|
|
|
for j := 0; j < rows; j++ {
|
|
var up, down float64
|
|
switch j {
|
|
case 0:
|
|
if rows == 1 {
|
|
up = 0.5
|
|
} else {
|
|
up = (h.GridXYZ.Y(1) - h.GridXYZ.Y(0)) / 2
|
|
}
|
|
down = -up
|
|
case rows - 1:
|
|
up = (h.GridXYZ.Y(rows-1) - h.GridXYZ.Y(rows-2)) / 2
|
|
down = -up
|
|
default:
|
|
up = (h.GridXYZ.Y(j+1) - h.GridXYZ.Y(j)) / 2
|
|
down = -(h.GridXYZ.Y(j) - h.GridXYZ.Y(j-1)) / 2
|
|
}
|
|
|
|
x, y := trX(h.GridXYZ.X(i)+left), trY(h.GridXYZ.Y(j)+down)
|
|
dx, dy := trX(h.GridXYZ.X(i)+right), trY(h.GridXYZ.Y(j)+up)
|
|
|
|
if !c.Contains(vg.Point{X: x, Y: y}) || !c.Contains(vg.Point{X: dx, Y: dy}) {
|
|
continue
|
|
}
|
|
|
|
pa = pa[:0]
|
|
pa.Move(vg.Point{X: x, Y: y})
|
|
pa.Line(vg.Point{X: dx, Y: y})
|
|
pa.Line(vg.Point{X: dx, Y: dy})
|
|
pa.Line(vg.Point{X: x, Y: dy})
|
|
pa.Close()
|
|
|
|
var col color.Color
|
|
switch v := h.GridXYZ.Z(i, j); {
|
|
case v < h.Min:
|
|
col = h.Underflow
|
|
case v > h.Max:
|
|
col = h.Overflow
|
|
case math.IsNaN(v), math.IsInf(ps, 0):
|
|
col = h.NaN
|
|
default:
|
|
col = pal[int((v-h.Min)*ps+0.5)] // Apply palette scaling.
|
|
}
|
|
if col != nil {
|
|
c.SetColor(col)
|
|
c.Fill(pa)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// DataRange implements the DataRange method
|
|
// of the plot.DataRanger interface.
|
|
func (h *HeatMap) DataRange() (xmin, xmax, ymin, ymax float64) {
|
|
c, r := h.GridXYZ.Dims()
|
|
switch c {
|
|
case 1: // Make a unit length when there is no neighbour.
|
|
xmax = h.GridXYZ.X(0) + 0.5
|
|
xmin = h.GridXYZ.X(0) - 0.5
|
|
default:
|
|
xmax = h.GridXYZ.X(c-1) + (h.GridXYZ.X(c-1)-h.GridXYZ.X(c-2))/2
|
|
xmin = h.GridXYZ.X(0) - (h.GridXYZ.X(1)-h.GridXYZ.X(0))/2
|
|
}
|
|
switch r {
|
|
case 1: // Make a unit length when there is no neighbour.
|
|
ymax = h.GridXYZ.Y(0) + 0.5
|
|
ymin = h.GridXYZ.Y(0) - 0.5
|
|
default:
|
|
ymax = h.GridXYZ.Y(r-1) + (h.GridXYZ.Y(r-1)-h.GridXYZ.Y(r-2))/2
|
|
ymin = h.GridXYZ.Y(0) - (h.GridXYZ.Y(1)-h.GridXYZ.Y(0))/2
|
|
}
|
|
return xmin, xmax, ymin, ymax
|
|
}
|
|
|
|
// GlyphBoxes implements the GlyphBoxes method
|
|
// of the plot.GlyphBoxer interface.
|
|
func (h *HeatMap) GlyphBoxes(plt *plot.Plot) []plot.GlyphBox {
|
|
c, r := h.GridXYZ.Dims()
|
|
b := make([]plot.GlyphBox, 0, r*c)
|
|
for i := 0; i < c; i++ {
|
|
for j := 0; j < r; j++ {
|
|
b = append(b, plot.GlyphBox{
|
|
X: plt.X.Norm(h.GridXYZ.X(i)),
|
|
Y: plt.Y.Norm(h.GridXYZ.Y(j)),
|
|
Rectangle: vg.Rectangle{
|
|
Min: vg.Point{X: -5, Y: -5},
|
|
Max: vg.Point{X: +5, Y: +5},
|
|
},
|
|
})
|
|
}
|
|
}
|
|
return b
|
|
}
|