plotter: Add ColorBar legend for palette.ColorMap

plotter/colorbar.go

@@ -0,0 +1,107 @@
+// Copyright ©2017 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"
+
+	"github.com/gonum/plot"
+	"github.com/gonum/plot/palette"
+	"github.com/gonum/plot/vg"
+	"github.com/gonum/plot/vg/draw"
+)
+
+// ColorBar is a plot.Plotter that draws a color bar legend for a ColorMap.
+type ColorBar struct {
+	// Vertical determines wether the legend will be
+	// plotted vertically or horizontally.
+	// The default is false (horizontal).
+	Vertical bool
+
+	// Colors specifies the number of colors to be
+	// shown in the legend. The default is 255.
+	Colors int
+
+	cm palette.ColorMap
+}
+
+// NewColorBar creates a new ColorBar plotter.
+func NewColorBar(cm palette.ColorMap) *ColorBar {
+	return &ColorBar{
+		cm:     cm,
+		Colors: 255,
+	}
+}
+
+// Plot implements the Plot method of the plot.Plotter interface.
+func (l *ColorBar) Plot(c draw.Canvas, p *plot.Plot) {
+	if l.cm.Max() == l.cm.Min() {
+		panic("palette: ColorMap Max==Min")
+	}
+	var img *image.NRGBA64
+	var xmin, xmax, ymin, ymax vg.Length
+	if l.Vertical {
+		trX, trY := p.Transforms(&c)
+		xmin = trX(l.cm.Min())
+		ymin = trY(0)
+		xmax = trX(l.cm.Max())
+		ymax = trY(1)
+		img = image.NewNRGBA64(image.Rectangle{
+			Min: image.Point{X: 0, Y: 0},
+			Max: image.Point{X: 1, Y: l.Colors},
+		})
+		for i := 0; i < l.Colors; i++ {
+			color, err := l.cm.At(float64(i) / float64(l.Colors-1))
+			if err != nil {
+				panic(err)
+			}
+			if l.Vertical {
+				img.Set(0, l.Colors-1-i, color)
+			} else {
+				img.Set(0, i, color)
+			}
+		}
+	} else {
+		trX, trY := p.Transforms(&c)
+		ymin = trY(l.cm.Min())
+		xmin = trX(0)
+		ymax = trY(l.cm.Max())
+		xmax = trX(1)
+		img = image.NewNRGBA64(image.Rectangle{
+			Min: image.Point{X: 0, Y: 0},
+			Max: image.Point{X: l.Colors, Y: 1},
+		})
+		for i := 0; i < l.Colors; i++ {
+			color, err := l.cm.At(float64(i) / float64(l.Colors-1))
+			if err != nil {
+				panic(err)
+			}
+			img.Set(i, 0, color)
+		}
+	}
+	rect := vg.Rectangle{
+		Min: vg.Point{X: xmin, Y: ymin},
+		Max: vg.Point{X: xmax, Y: ymax},
+	}
+	c.DrawImage(rect, img)
+}
+
+// DataRange implements the DataRange method
+// of the plot.DataRanger interface.
+func (l *ColorBar) DataRange() (xmin, xmax, ymin, ymax float64) {
+	if l.cm.Max() == l.cm.Min() {
+		panic("palette: ColorMap Max==Min")
+	}
+	if l.Vertical {
+		return 0, 1, l.cm.Min(), l.cm.Max()
+	}
+	return l.cm.Min(), l.cm.Max(), 0, 1
+}
+
+// GlyphBoxes implements the GlyphBoxes method
+// of the plot.GlyphBoxer interface.
+func (l *ColorBar) GlyphBoxes(plt *plot.Plot) []plot.GlyphBox {
+	return nil
+}

plotter/colorbar_test.go

@@ -0,0 +1,59 @@
+// Copyright ©2017 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 (
+	"log"
+	"testing"
+
+	"github.com/gonum/plot"
+	"github.com/gonum/plot/internal/cmpimg"
+	"github.com/gonum/plot/palette/moreland"
+)
+
+func ExampleColorBar_horizontal() {
+	p, err := plot.New()
+	if err != nil {
+		log.Panic(err)
+	}
+	cm := moreland.ExtendedBlackBody()
+	cm.SetMax(1)
+	l := NewColorBar(cm)
+	p.Add(l)
+	p.HideY()
+	p.X.Padding = 0
+	p.Title.Text = "Title"
+
+	if err = p.Save(300, 48, "testdata/colorBarHorizontal.png"); err != nil {
+		log.Panic(err)
+	}
+}
+
+func TestColorBar_horizontal(t *testing.T) {
+	cmpimg.CheckPlot(ExampleColorBar_horizontal, t, "colorBarHorizontal.png")
+}
+
+func ExampleColorBar_vertical() {
+	p, err := plot.New()
+	if err != nil {
+		log.Panic(err)
+	}
+	cm := moreland.ExtendedBlackBody()
+	cm.SetMax(1)
+	l := NewColorBar(cm)
+	l.Vertical = true
+	p.Add(l)
+	p.HideX()
+	p.Y.Padding = 0
+	p.Title.Text = "Title"
+
+	if err = p.Save(40, 300, "testdata/colorBarVertical.png"); err != nil {
+		log.Panic(err)
+	}
+}
+
+func TestColorBar_vertical(t *testing.T) {
+	cmpimg.CheckPlot(ExampleColorBar_vertical, t, "colorBarVertical.png")
+}