hplot: add automatic legend creation to H2D

Signed-off-by: Sebastien Binet <binet@cern.ch>

hplot/h2d.go

@@ -5,6 +5,8 @@
 package hplot
 
 import (
+	"fmt"
+
 	"go-hep.org/x/hep/hbook"
 	"gonum.org/v1/plot"
 	"gonum.org/v1/plot/palette"
@@ -58,6 +60,30 @@ func (h *H2D) GlyphBoxes(p *plot.Plot) []plot.GlyphBox {
 	return h.HeatMap.GlyphBoxes(p)
 }
 
+// Legend returns a legend constructed from the 2-dim data and palette.
+func (h *H2D) Legend() Legend {
+	legend := NewLegend()
+	thumbs := plotter.PaletteThumbnailers(h.HeatMap.Palette)
+	for i := len(thumbs) - 1; i >= 0; i-- {
+		t := thumbs[i]
+		if i != 0 && i != len(thumbs)-1 {
+			legend.Add("", t)
+			continue
+		}
+		var val float64
+		switch i {
+		case 0:
+			val = h.HeatMap.Min
+		case len(thumbs) - 1:
+			val = h.HeatMap.Max
+		}
+		legend.Add(fmt.Sprintf("%.2g", val), t)
+	}
+	legend.Top = true
+
+	return legend
+}
+
 // check interfaces
 var _ plot.Plotter = (*H2D)(nil)
 var _ plot.DataRanger = (*H2D)(nil)

hplot/h2d_example_test.go

@@ -50,3 +50,41 @@ func ExampleH2D() {
 		log.Fatal(err)
 	}
 }
+
+func ExampleH2D_withLegend() {
+	h2d := hbook.NewH2D(100, -10, 10, 100, -10, 10)
+
+	const npoints = 10000
+
+	dist, ok := distmv.NewNormal(
+		[]float64{0, 1},
+		mat.NewSymDense(2, []float64{4, 0, 0, 2}),
+		rand.New(rand.NewSource(1234)),
+	)
+	if !ok {
+		log.Fatalf("error creating distmv.Normal")
+	}
+
+	v := make([]float64, 2)
+	// Draw some random values from the standard
+	// normal distribution.
+	for i := 0; i < npoints; i++ {
+		v = dist.Rand(v)
+		h2d.Fill(v[0], v[1], 1)
+	}
+	h := hplot.NewH2D(h2d, nil)
+
+	p := hplot.New()
+	p.Title.Text = "Hist-2D"
+	p.X.Label.Text = "x"
+	p.Y.Label.Text = "y"
+
+	p.Add(h)
+	p.Add(plotter.NewGrid())
+
+	fig := hplot.Figure(p, hplot.WithLegend(h.Legend()))
+	err := hplot.Save(fig, 10*vg.Centimeter, 10*vg.Centimeter, "testdata/h2d_plot_legend.png")
+	if err != nil {
+		log.Fatal(err)
+	}
+}

hplot/h2d_test.go

@@ -19,6 +19,10 @@ func TestH2D(t *testing.T) {
 	checkPlot(cmpimg.CheckPlot)(ExampleH2D, t, "h2d_plot.png")
 }
 
+func TestH2DWithLegend(t *testing.T) {
+	checkPlot(cmpimg.CheckPlot)(ExampleH2D_withLegend, t, "h2d_plot_legend.png")
+}
+
 func TestH2DABCD(t *testing.T) {
 	checkPlot(cmpimg.CheckPlot)(func() {
 		h := hbook.NewH2D(2, 0, 2, 2, 0, 2)