Merge e68ae8f into 4caad36

plotter/field.go

@@ -38,18 +38,24 @@ type Field struct {
 
 	// DrawGlyph is the user hook to draw a field
 	// vector glyph. The function should draw a unit
-	// vector to (1, 0) on the vg.Canvas, c.
+	// vector to (1, 0) on the vg.Canvas, c with the
+	// sty LineStyle. The Field plotter will rotate
+	// and scale the unit vector appropriately.
+	// If the magnitude of v is zero, no scaling or
+	// rotation is performed.
+	//
 	// The direction and magnitude of v can be used
 	// to determine properties of the glyph drawing
 	// but should not be used to determine size or
 	// directions of the glyph.
 	//
 	// If DrawGlyph is nil, a simple arrow will be
 	// drawn.
-	DrawGlyph func(c vg.Canvas, v XY)
+	DrawGlyph func(c vg.Canvas, sty draw.LineStyle, v XY)
 
 	// LineStyle is the style of the line used to
-	// render vectors if DrawGlyph is nil.
+	// render vectors when DrawGlyph is nil.
+	// Otherwise it is passed to DrawGlyph.
 	LineStyle draw.LineStyle
 
 	// max define the dynamic range of the field.
@@ -124,16 +130,22 @@ func (f *Field) Plot(c draw.Canvas, plt *plot.Plot) {
 
 			c.Push()
 			c.Translate(vg.Point{X: (x + dx) / 2, Y: (y + dy) / 2})
+
 			v := f.FieldXY.Vector(i, j)
-			c.Rotate(math.Atan2(v.Y, v.X))
 			s := math.Hypot(v.X, v.Y) / (2 * f.max)
-			v.X *= s
-			v.Y *= s
-			c.Scale(s*float64(dx-x), s*float64(dy-y))
+			// Do not scale when the vector is zero, otherwise the
+			// user cannot render special-case glyphs for that case.
+			if s != 0 {
+				c.Rotate(math.Atan2(v.Y, v.X))
+				c.Scale(s*float64(dx-x), s*float64(dy-y))
+			}
+			v.X /= f.max
+			v.Y /= f.max
+
 			if f.DrawGlyph == nil {
 				drawVector(c, v)
 			} else {
-				f.DrawGlyph(c, v)
+				f.DrawGlyph(c, f.LineStyle, v)
 			}
 			c.Pop()
 		}

plotter/field_test.go

@@ -5,13 +5,16 @@
 package plotter_test
 
 import (
+	"image/color"
 	"image/png"
 	"log"
+	"math"
 	"os"
 	"testing"
 
 	"gonum.org/v1/plot"
 	"gonum.org/v1/plot/cmpimg"
+	"gonum.org/v1/plot/palette/moreland"
 	"gonum.org/v1/plot/plotter"
 	"gonum.org/v1/plot/vg"
 	"gonum.org/v1/plot/vg/draw"
@@ -79,6 +82,78 @@ func TestField(t *testing.T) {
 	cmpimg.CheckPlot(ExampleField, t, "field.png")
 }
 
+func ExampleField_colors() {
+	f := plotter.NewField(field{
+		r: 5, c: 9,
+		fn: func(x, y float64) plotter.XY {
+			return plotter.XY{
+				X: -0.75*x + y,
+				Y: -0.75*y - x,
+			}
+		},
+	})
+
+	pal := moreland.ExtendedBlackBody()
+	pal.SetMin(0)
+	pal.SetMax(1.1) // Use 1.1 to make highest magnitude vectors visible on white.
+
+	// Provide a DrawGlyph function to render a custom
+	// vector instead of the default monochrome arrow.
+	f.DrawGlyph = func(c vg.Canvas, sty draw.LineStyle, v plotter.XY) {
+		c.Push()
+		defer c.Pop()
+		mag := math.Hypot(v.X, v.Y)
+		var pa vg.Path
+		if mag == 0 {
+			// Draw a black dot for zero vectors.
+			c.SetColor(color.Black)
+			pa.Move(vg.Point{X: sty.Width})
+			pa.Arc(vg.Point{}, sty.Width, 0, 2*math.Pi)
+			pa.Close()
+			c.Fill(pa)
+			return
+		}
+		// Choose a color from the palette for the magnitude.
+		col, err := pal.At(mag)
+		if err != nil {
+			panic(err)
+		}
+		c.SetColor(col)
+		pa.Move(vg.Point{})
+		pa.Line(vg.Point{X: 1, Y: 0})
+		pa.Close()
+		c.Stroke(pa)
+	}
+
+	p, err := plot.New()
+	if err != nil {
+		log.Panic(err)
+	}
+	p.Title.Text = "Vortex"
+
+	p.X.Tick.Marker = integerTicks{}
+	p.Y.Tick.Marker = integerTicks{}
+
+	p.Add(f)
+
+	img := vgimg.New(250, 175)
+	dc := draw.New(img)
+
+	p.Draw(dc)
+	w, err := os.Create("testdata/color_field.png")
+	if err != nil {
+		log.Panic(err)
+	}
+	png := vgimg.PngCanvas{Canvas: img}
+	if _, err = png.WriteTo(w); err != nil {
+		log.Panic(err)
+	}
+}
+
+func TestFieldColors(t *testing.T) {
+	cmpimg.CheckPlot(ExampleField_colors, t, "color_field.png")
+}
+
 func ExampleField_gophers() {
 	file, err := os.Open("testdata/gopher_running.png")
 	if err != nil {
@@ -102,7 +177,12 @@ func ExampleField_gophers() {
 
 	// Provide a DrawGlyph function to render a custom
 	// vector glyph instead of the default arrow.
-	f.DrawGlyph = func(c vg.Canvas, v plotter.XY) {
+	f.DrawGlyph = func(c vg.Canvas, _ draw.LineStyle, v plotter.XY) {
+		// The canvas is unscaled if the vector has a zero
+		// magnitude, so return in that case.
+		if math.Hypot(v.X, v.Y) == 0 {
+			return
+		}
 		// Vector glyphs are scaled to half unit length by the
 		// plotter, so scale the gopher to twice unit size so
 		// it fits the cell, and center on the cell.
@@ -135,13 +215,6 @@ func ExampleField_gophers() {
 	}
 }
 
-func max(a, b int) int {
-	if a > b {
-		return a
-	}
-	return b
-}
-
 func TestFieldGophers(t *testing.T) {
 	cmpimg.CheckPlot(ExampleField_gophers, t, "gopher_field.png")
 }