heat.go

Documentation: gonum.org/v1/plot/plotter

     1  // Copyright ©2015 The Gonum Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package plotter
     6  
     7  import (
     8  	"image"
     9  	"image/color"
    10  	"math"
    11  
    12  	"gonum.org/v1/plot"
    13  	"gonum.org/v1/plot/palette"
    14  	"gonum.org/v1/plot/vg"
    15  	"gonum.org/v1/plot/vg/draw"
    16  )
    17  
    18  // GridXYZ describes three dimensional data where the X and Y
    19  // coordinates are arranged on a rectangular grid.
    20  type GridXYZ interface {
    21  	// Dims returns the dimensions of the grid.
    22  	Dims() (c, r int)
    23  
    24  	// Z returns the value of a grid value at (c, r).
    25  	// It will panic if c or r are out of bounds for the grid.
    26  	Z(c, r int) float64
    27  
    28  	// X returns the coordinate for the column at the index c.
    29  	// It will panic if c is out of bounds for the grid.
    30  	X(c int) float64
    31  
    32  	// Y returns the coordinate for the row at the index r.
    33  	// It will panic if r is out of bounds for the grid.
    34  	Y(r int) float64
    35  }
    36  
    37  // HeatMap implements the Plotter interface, drawing
    38  // a heat map of the values in the GridXYZ field.
    39  type HeatMap struct {
    40  	GridXYZ GridXYZ
    41  
    42  	// Palette is the color palette used to render
    43  	// the heat map. Palette must not be nil or
    44  	// return a zero length []color.Color.
    45  	Palette palette.Palette
    46  
    47  	// Underflow and Overflow are colors used to fill
    48  	// heat map elements outside the dynamic range
    49  	// defined by Min and Max.
    50  	Underflow color.Color
    51  	Overflow  color.Color
    52  
    53  	// NaN is the color used to fill heat map elements
    54  	// that are NaN or do not map to a unique palette
    55  	// color.
    56  	NaN color.Color
    57  
    58  	// Min and Max define the dynamic range of the
    59  	// heat map.
    60  	Min, Max float64
    61  
    62  	// Rasterized indicates whether the heatmap
    63  	// should be produced using raster-based drawing.
    64  	Rasterized bool
    65  }
    66  
    67  // NewHeatMap creates as new heat map plotter for the given data,
    68  // using the provided palette. If g has Min and Max methods that return
    69  // a float, those returned values are used to set the respective HeatMap
    70  // fields. If the returned HeatMap is used when Min is greater than Max,
    71  // the Plot method will panic.
    72  func NewHeatMap(g GridXYZ, p palette.Palette) *HeatMap {
    73  	var min, max float64
    74  	type minMaxer interface {
    75  		Min() float64
    76  		Max() float64
    77  	}
    78  	switch g := g.(type) {
    79  	case minMaxer:
    80  		min, max = g.Min(), g.Max()
    81  	default:
    82  		min, max = math.Inf(1), math.Inf(-1)
    83  		c, r := g.Dims()
    84  		for i := 0; i < c; i++ {
    85  			for j := 0; j < r; j++ {
    86  				v := g.Z(i, j)
    87  				if math.IsNaN(v) {
    88  					continue
    89  				}
    90  				min = math.Min(min, v)
    91  				max = math.Max(max, v)
    92  			}
    93  		}
    94  	}
    95  
    96  	return &HeatMap{
    97  		GridXYZ: g,
    98  		Palette: p,
    99  		Min:     min,
   100  		Max:     max,
   101  	}
   102  }
   103  
   104  // Plot implements the Plot method of the plot.Plotter interface.
   105  func (h *HeatMap) Plot(c draw.Canvas, plt *plot.Plot) {
   106  	if h.Rasterized {
   107  		h.plotRasterized(c, plt)
   108  	} else {
   109  		h.plotVectorized(c, plt)
   110  	}
   111  }
   112  
   113  // plotRasterized plots the heatmap using raster-based drawing.
   114  func (h *HeatMap) plotRasterized(c draw.Canvas, plt *plot.Plot) {
   115  	cols, rows := h.GridXYZ.Dims()
   116  	img := image.NewRGBA64(image.Rectangle{
   117  		Min: image.Point{X: 0, Y: 0},
   118  		Max: image.Point{X: cols, Y: rows},
   119  	})
   120  
   121  	pal := h.Palette.Colors()
   122  	ps := float64(len(pal)-1) / (h.Max - h.Min)
   123  	for i := 0; i < cols; i++ {
   124  		for j := 0; j < rows; j++ {
   125  			var col color.Color
   126  			switch v := h.GridXYZ.Z(i, j); {
   127  			case v < h.Min:
   128  				col = h.Underflow
   129  			case v > h.Max:
   130  				col = h.Overflow
   131  			case math.IsNaN(v), math.IsInf(ps, 0):
   132  				col = h.NaN
   133  			default:
   134  				col = pal[int((v-h.Min)*ps+0.5)] // Apply palette scaling.
   135  			}
   136  
   137  			if col != nil {
   138  				img.Set(i, rows-j-1, col)
   139  			}
   140  		}
   141  	}
   142  
   143  	xmin, xmax, ymin, ymax := h.DataRange()
   144  	pImg := NewImage(img, xmin, ymin, xmax, ymax)
   145  	pImg.Plot(c, plt)
   146  }
   147  
   148  // plotVectorized plots the heatmap using vector-based drawing.
   149  func (h *HeatMap) plotVectorized(c draw.Canvas, plt *plot.Plot) {
   150  	if h.Min > h.Max {
   151  		panic("contour: invalid Z range: min greater than max")
   152  	}
   153  	pal := h.Palette.Colors()
   154  	if len(pal) == 0 {
   155  		panic("heatmap: empty palette")
   156  	}
   157  	// ps scales the palette uniformly across the data range.
   158  	ps := float64(len(pal)-1) / (h.Max - h.Min)
   159  
   160  	trX, trY := plt.Transforms(&c)
   161  
   162  	var pa vg.Path
   163  	cols, rows := h.GridXYZ.Dims()
   164  	for i := 0; i < cols; i++ {
   165  		var right, left float64
   166  		switch i {
   167  		case 0:
   168  			if cols == 1 {
   169  				right = 0.5
   170  			} else {
   171  				right = (h.GridXYZ.X(1) - h.GridXYZ.X(0)) / 2
   172  			}
   173  			left = -right
   174  		case cols - 1:
   175  			right = (h.GridXYZ.X(cols-1) - h.GridXYZ.X(cols-2)) / 2
   176  			left = -right
   177  		default:
   178  			right = (h.GridXYZ.X(i+1) - h.GridXYZ.X(i)) / 2
   179  			left = -(h.GridXYZ.X(i) - h.GridXYZ.X(i-1)) / 2
   180  		}
   181  
   182  		for j := 0; j < rows; j++ {
   183  			var up, down float64
   184  			switch j {
   185  			case 0:
   186  				if rows == 1 {
   187  					up = 0.5
   188  				} else {
   189  					up = (h.GridXYZ.Y(1) - h.GridXYZ.Y(0)) / 2
   190  				}
   191  				down = -up
   192  			case rows - 1:
   193  				up = (h.GridXYZ.Y(rows-1) - h.GridXYZ.Y(rows-2)) / 2
   194  				down = -up
   195  			default:
   196  				up = (h.GridXYZ.Y(j+1) - h.GridXYZ.Y(j)) / 2
   197  				down = -(h.GridXYZ.Y(j) - h.GridXYZ.Y(j-1)) / 2
   198  			}
   199  
   200  			x, y := trX(h.GridXYZ.X(i)+left), trY(h.GridXYZ.Y(j)+down)
   201  			dx, dy := trX(h.GridXYZ.X(i)+right), trY(h.GridXYZ.Y(j)+up)
   202  
   203  			if !c.Contains(vg.Point{X: x, Y: y}) || !c.Contains(vg.Point{X: dx, Y: dy}) {
   204  				continue
   205  			}
   206  
   207  			pa = pa[:0]
   208  			pa.Move(vg.Point{X: x, Y: y})
   209  			pa.Line(vg.Point{X: dx, Y: y})
   210  			pa.Line(vg.Point{X: dx, Y: dy})
   211  			pa.Line(vg.Point{X: x, Y: dy})
   212  			pa.Close()
   213  
   214  			var col color.Color
   215  			switch v := h.GridXYZ.Z(i, j); {
   216  			case v < h.Min:
   217  				col = h.Underflow
   218  			case v > h.Max:
   219  				col = h.Overflow
   220  			case math.IsNaN(v), math.IsInf(ps, 0):
   221  				col = h.NaN
   222  			default:
   223  				col = pal[int((v-h.Min)*ps+0.5)] // Apply palette scaling.
   224  			}
   225  			if col != nil {
   226  				c.SetColor(col)
   227  				c.Fill(pa)
   228  			}
   229  		}
   230  	}
   231  }
   232  
   233  // DataRange implements the DataRange method
   234  // of the plot.DataRanger interface.
   235  func (h *HeatMap) DataRange() (xmin, xmax, ymin, ymax float64) {
   236  	c, r := h.GridXYZ.Dims()
   237  	switch c {
   238  	case 1: // Make a unit length when there is no neighbour.
   239  		xmax = h.GridXYZ.X(0) + 0.5
   240  		xmin = h.GridXYZ.X(0) - 0.5
   241  	default:
   242  		xmax = h.GridXYZ.X(c-1) + (h.GridXYZ.X(c-1)-h.GridXYZ.X(c-2))/2
   243  		xmin = h.GridXYZ.X(0) - (h.GridXYZ.X(1)-h.GridXYZ.X(0))/2
   244  	}
   245  	switch r {
   246  	case 1: // Make a unit length when there is no neighbour.
   247  		ymax = h.GridXYZ.Y(0) + 0.5
   248  		ymin = h.GridXYZ.Y(0) - 0.5
   249  	default:
   250  		ymax = h.GridXYZ.Y(r-1) + (h.GridXYZ.Y(r-1)-h.GridXYZ.Y(r-2))/2
   251  		ymin = h.GridXYZ.Y(0) - (h.GridXYZ.Y(1)-h.GridXYZ.Y(0))/2
   252  	}
   253  	return xmin, xmax, ymin, ymax
   254  }
   255  
   256  // GlyphBoxes implements the GlyphBoxes method
   257  // of the plot.GlyphBoxer interface.
   258  func (h *HeatMap) GlyphBoxes(plt *plot.Plot) []plot.GlyphBox {
   259  	c, r := h.GridXYZ.Dims()
   260  	b := make([]plot.GlyphBox, 0, r*c)
   261  	for i := 0; i < c; i++ {
   262  		for j := 0; j < r; j++ {
   263  			b = append(b, plot.GlyphBox{
   264  				X: plt.X.Norm(h.GridXYZ.X(i)),
   265  				Y: plt.Y.Norm(h.GridXYZ.Y(j)),
   266  				Rectangle: vg.Rectangle{
   267  					Min: vg.Point{X: -5, Y: -5},
   268  					Max: vg.Point{X: +5, Y: +5},
   269  				},
   270  			})
   271  		}
   272  	}
   273  	return b
   274  }
   275
View as plain text