sankey_example_test.go

Documentation: gonum.org/v1/plot/plotter

     1  // Copyright ©2016 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_test
     6  
     7  import (
     8  	"fmt"
     9  	"image/color"
    10  	"log"
    11  	"os"
    12  
    13  	"gonum.org/v1/plot"
    14  	"gonum.org/v1/plot/plotter"
    15  	"gonum.org/v1/plot/text"
    16  	"gonum.org/v1/plot/vg"
    17  	"gonum.org/v1/plot/vg/draw"
    18  	"gonum.org/v1/plot/vg/vgimg"
    19  )
    20  
    21  // ExampleSankey_sample creates a simple sankey diagram.
    22  // The output can be found at https://github.com/gonum/plot/blob/master/plotter/testdata/sankeySimple_golden.png.
    23  func ExampleSankey_simple() {
    24  	p := plot.New()
    25  
    26  	// Define the stock categories
    27  	const (
    28  		treeType int = iota
    29  		consumer
    30  		fate
    31  	)
    32  	categoryLabels := []string{"Tree type", "Consumer", "Fate"}
    33  
    34  	flows := []plotter.Flow{
    35  		{
    36  			SourceCategory:   treeType,
    37  			SourceLabel:      "Large",
    38  			ReceptorCategory: consumer,
    39  			ReceptorLabel:    "Mohamed",
    40  			Value:            5,
    41  		},
    42  		{
    43  			SourceCategory:   treeType,
    44  			SourceLabel:      "Small",
    45  			ReceptorCategory: consumer,
    46  			ReceptorLabel:    "Mohamed",
    47  			Value:            2,
    48  		},
    49  		{
    50  			SourceCategory:   treeType,
    51  			SourceLabel:      "Large",
    52  			ReceptorCategory: consumer,
    53  			ReceptorLabel:    "Sofia",
    54  			Value:            3,
    55  		},
    56  		{
    57  			SourceCategory:   treeType,
    58  			SourceLabel:      "Small",
    59  			ReceptorCategory: consumer,
    60  			ReceptorLabel:    "Sofia",
    61  			Value:            1,
    62  		},
    63  		{
    64  			SourceCategory:   treeType,
    65  			SourceLabel:      "Large",
    66  			ReceptorCategory: consumer,
    67  			ReceptorLabel:    "Wei",
    68  			Value:            6,
    69  		},
    70  		{
    71  			SourceCategory:   consumer,
    72  			SourceLabel:      "Mohamed",
    73  			ReceptorCategory: fate,
    74  			ReceptorLabel:    "Eaten",
    75  			Value:            6,
    76  		},
    77  		{
    78  			SourceCategory:   consumer,
    79  			SourceLabel:      "Mohamed",
    80  			ReceptorCategory: fate,
    81  			ReceptorLabel:    "Waste",
    82  			Value:            1,
    83  		},
    84  		{
    85  			SourceCategory:   consumer,
    86  			SourceLabel:      "Sofia",
    87  			ReceptorCategory: fate,
    88  			ReceptorLabel:    "Eaten",
    89  			Value:            3,
    90  		},
    91  		{
    92  			SourceCategory:   consumer,
    93  			SourceLabel:      "Sofia",
    94  			ReceptorCategory: fate,
    95  			ReceptorLabel:    "Waste",
    96  			Value:            0.5, // An unbalanced flow
    97  		},
    98  		{
    99  			SourceCategory:   consumer,
   100  			SourceLabel:      "Wei",
   101  			ReceptorCategory: fate,
   102  			ReceptorLabel:    "Eaten",
   103  			Value:            5,
   104  		},
   105  		{
   106  			SourceCategory:   consumer,
   107  			SourceLabel:      "Wei",
   108  			ReceptorCategory: fate,
   109  			ReceptorLabel:    "Waste",
   110  			Value:            1,
   111  		},
   112  		{
   113  			SourceCategory:   treeType,
   114  			SourceLabel:      "Large",
   115  			ReceptorCategory: fate,
   116  			ReceptorLabel:    "Waste",
   117  			Value:            1,
   118  		},
   119  		{
   120  			SourceCategory:   treeType,
   121  			SourceLabel:      "Small",
   122  			ReceptorCategory: fate,
   123  			ReceptorLabel:    "Waste",
   124  			Value:            0.3,
   125  		},
   126  	}
   127  
   128  	sankey, err := plotter.NewSankey(flows...)
   129  	if err != nil {
   130  		log.Panic(err)
   131  	}
   132  	p.Add(sankey)
   133  	p.Y.Label.Text = "Number of apples"
   134  	p.NominalX(categoryLabels...)
   135  	err = p.Save(vg.Points(300), vg.Points(180), "testdata/sankeySimple.png")
   136  	if err != nil {
   137  		log.Panic(err)
   138  	}
   139  }
   140  
   141  // ExampleSankey_grouped creates a sankey diagram with grouped flows.
   142  // The output can be found at https://github.com/gonum/plot/blob/master/plotter/testdata/sankeyGrouped_golden.png.
   143  func ExampleSankey_grouped() {
   144  	p := plot.New()
   145  	c := vgimg.New(vg.Points(300), vg.Points(180))
   146  	dc := draw.New(c)
   147  
   148  	// Define the stock categories
   149  	const (
   150  		treeType int = iota
   151  		consumer
   152  		fate
   153  	)
   154  	categoryLabels := []string{"Tree type", "Consumer", "Fate"}
   155  
   156  	flows := []plotter.Flow{
   157  		{
   158  			SourceCategory:   treeType,
   159  			SourceLabel:      "LargeLargeLargeLargeLargeLargeLargeLargeLarge",
   160  			ReceptorCategory: consumer,
   161  			ReceptorLabel:    "Mohamed",
   162  			Group:            "Apples",
   163  			Value:            5,
   164  		},
   165  		{
   166  			SourceCategory:   treeType,
   167  			SourceLabel:      "LargeLargeLargeLargeLargeLargeLargeLargeLarge",
   168  			ReceptorCategory: consumer,
   169  			ReceptorLabel:    "Mohamed",
   170  			Group:            "Dates",
   171  			Value:            3,
   172  		},
   173  		{
   174  			SourceCategory:   treeType,
   175  			SourceLabel:      "Small",
   176  			ReceptorCategory: consumer,
   177  			ReceptorLabel:    "Mohamed",
   178  			Group:            "Lychees",
   179  			Value:            2,
   180  		},
   181  		{
   182  			SourceCategory:   treeType,
   183  			SourceLabel:      "LargeLargeLargeLargeLargeLargeLargeLargeLarge",
   184  			ReceptorCategory: consumer,
   185  			ReceptorLabel:    "Sofia",
   186  			Group:            "Apples",
   187  			Value:            3,
   188  		},
   189  		{
   190  			SourceCategory:   treeType,
   191  			SourceLabel:      "LargeLargeLargeLargeLargeLargeLargeLargeLarge",
   192  			ReceptorCategory: consumer,
   193  			ReceptorLabel:    "Sofia",
   194  			Group:            "Dates",
   195  			Value:            4,
   196  		},
   197  		{
   198  			SourceCategory:   treeType,
   199  			SourceLabel:      "Small",
   200  			ReceptorCategory: consumer,
   201  			ReceptorLabel:    "Sofia",
   202  			Group:            "Apples",
   203  			Value:            1,
   204  		},
   205  		{
   206  			SourceCategory:   treeType,
   207  			SourceLabel:      "LargeLargeLargeLargeLargeLargeLargeLargeLarge",
   208  			ReceptorCategory: consumer,
   209  			ReceptorLabel:    "Wei",
   210  			Group:            "Lychees",
   211  			Value:            6,
   212  		},
   213  		{
   214  			SourceCategory:   treeType,
   215  			SourceLabel:      "Small",
   216  			ReceptorCategory: consumer,
   217  			ReceptorLabel:    "Wei",
   218  			Group:            "Apples",
   219  			Value:            3,
   220  		},
   221  		{
   222  			SourceCategory:   consumer,
   223  			SourceLabel:      "Mohamed",
   224  			ReceptorCategory: fate,
   225  			ReceptorLabel:    "Eaten",
   226  			Group:            "Apples",
   227  			Value:            4,
   228  		},
   229  		{
   230  			SourceCategory:   consumer,
   231  			SourceLabel:      "Mohamed",
   232  			ReceptorCategory: fate,
   233  			ReceptorLabel:    "Waste",
   234  			Group:            "Apples",
   235  			Value:            1,
   236  		},
   237  		{
   238  			SourceCategory:   consumer,
   239  			SourceLabel:      "Mohamed",
   240  			ReceptorCategory: fate,
   241  			ReceptorLabel:    "Eaten",
   242  			Group:            "Dates",
   243  			Value:            3,
   244  		},
   245  		{
   246  			SourceCategory:   consumer,
   247  			SourceLabel:      "Mohamed",
   248  			ReceptorCategory: fate,
   249  			ReceptorLabel:    "Waste",
   250  			Group:            "Lychees",
   251  			Value:            2,
   252  		},
   253  		{
   254  			SourceCategory:   consumer,
   255  			SourceLabel:      "Sofia",
   256  			ReceptorCategory: fate,
   257  			ReceptorLabel:    "Eaten",
   258  			Group:            "Apples",
   259  			Value:            4,
   260  		},
   261  		{
   262  			SourceCategory:   consumer,
   263  			SourceLabel:      "Sofia",
   264  			ReceptorCategory: fate,
   265  			ReceptorLabel:    "Eaten",
   266  			Group:            "Dates",
   267  			Value:            3,
   268  		},
   269  		{
   270  			SourceCategory:   consumer,
   271  			SourceLabel:      "Sofia",
   272  			ReceptorCategory: fate,
   273  			ReceptorLabel:    "Waste",
   274  			Group:            "Dates",
   275  			Value:            1,
   276  		},
   277  		{
   278  			SourceCategory:   consumer,
   279  			SourceLabel:      "Wei",
   280  			ReceptorCategory: fate,
   281  			ReceptorLabel:    "Eaten",
   282  			Group:            "Lychees",
   283  			Value:            6,
   284  		},
   285  		{
   286  			SourceCategory:   consumer,
   287  			SourceLabel:      "Wei",
   288  			ReceptorCategory: fate,
   289  			ReceptorLabel:    "Eaten",
   290  			Group:            "Apples",
   291  			Value:            2,
   292  		},
   293  		{
   294  			SourceCategory:   consumer,
   295  			SourceLabel:      "Wei",
   296  			ReceptorCategory: fate,
   297  			ReceptorLabel:    "Waste",
   298  			Group:            "Apples",
   299  			Value:            1,
   300  		},
   301  		{
   302  			SourceCategory:   treeType,
   303  			SourceLabel:      "LargeLargeLargeLargeLargeLargeLargeLargeLarge",
   304  			ReceptorCategory: fate,
   305  			ReceptorLabel:    "Waste",
   306  			Group:            "Apples",
   307  			Value:            1,
   308  		},
   309  		{
   310  			SourceCategory:   treeType,
   311  			SourceLabel:      "LargeLargeLargeLargeLargeLargeLargeLargeLarge",
   312  			ReceptorCategory: fate,
   313  			ReceptorLabel:    "Waste",
   314  			Group:            "Dates",
   315  			Value:            1,
   316  		},
   317  		{
   318  			SourceCategory:   treeType,
   319  			SourceLabel:      "Small",
   320  			ReceptorCategory: fate,
   321  			ReceptorLabel:    "Waste",
   322  			Group:            "Lychees",
   323  			Value:            0.3,
   324  		},
   325  	}
   326  
   327  	sankey, err := plotter.NewSankey(flows...)
   328  	if err != nil {
   329  		log.Panic(err)
   330  	}
   331  
   332  	// Here we specify the FLowStyle function to set the
   333  	// colors of the different fruit groups.
   334  	sankey.FlowStyle = func(group string) (color.Color, draw.LineStyle) {
   335  		switch group {
   336  		case "Lychees":
   337  			return color.NRGBA{R: 242, G: 169, B: 178, A: 100}, sankey.LineStyle
   338  		case "Apples":
   339  			return color.NRGBA{R: 91, G: 194, B: 54, A: 100}, sankey.LineStyle
   340  		case "Dates":
   341  			return color.NRGBA{R: 112, G: 22, B: 0, A: 100}, sankey.LineStyle
   342  		default:
   343  			panic(fmt.Errorf("invalid group %s", group))
   344  		}
   345  	}
   346  
   347  	// Here we set the StockStyle function to give an example of
   348  	// setting a custom style for one of the stocks.
   349  	sankey.StockStyle = func(label string, category int) (string, text.Style, vg.Length, vg.Length, color.Color, draw.LineStyle) {
   350  		if label == "Small" && category == treeType {
   351  			// Here we demonstrate how to rotate the label text
   352  			// and change the style of the stock bar.
   353  			ts := sankey.TextStyle
   354  			ts.Rotation = 0.0
   355  			ts.XAlign = draw.XRight
   356  			ls := sankey.LineStyle
   357  			ls.Color = color.White
   358  			xOff := -sankey.StockBarWidth / 2
   359  			yOff := vg.Length(0)
   360  			return "small", ts, xOff, yOff, color.Black, ls
   361  		}
   362  		if label == "LargeLargeLargeLargeLargeLargeLargeLargeLarge" && category == treeType {
   363  			// Here we demonstrate how to replace a long label that doesn't fit
   364  			// in the existing space with a shorter version. Note that because
   365  			// we are not able to account for the difference between the overall
   366  			// canvas size and the size of the plotting area here, if a label
   367  			// was only slightly larger than the available space, it would not
   368  			// be caught and replaced.
   369  			min, max, err := sankey.StockRange(label, category)
   370  			if err != nil {
   371  				log.Panic(err)
   372  			}
   373  			_, yTr := p.Transforms(&dc)
   374  			barHeight := yTr(max) - yTr(min)
   375  			if sankey.TextStyle.Width(label) > barHeight {
   376  				return "large", sankey.TextStyle, 0, 0, color.White, sankey.LineStyle
   377  			}
   378  		}
   379  		return label, sankey.TextStyle, 0, 0, color.White, sankey.LineStyle
   380  	}
   381  
   382  	p.Add(sankey)
   383  	p.Y.Label.Text = "Number of fruit pieces"
   384  	p.NominalX(categoryLabels...)
   385  
   386  	legendLabels, thumbs := sankey.Thumbnailers()
   387  	for i, l := range legendLabels {
   388  		t := thumbs[i]
   389  		p.Legend.Add(l, t)
   390  	}
   391  	p.Legend.Top = true
   392  	p.X.Max = 3.05 // give room for the legend
   393  
   394  	// Add boundary boxes for debugging.
   395  	p.Add(plotter.NewGlyphBoxes())
   396  
   397  	p.Draw(dc)
   398  	pngimg := vgimg.PngCanvas{Canvas: c}
   399  	f, err := os.Create("testdata/sankeyGrouped.png")
   400  	if err != nil {
   401  		log.Panic(err)
   402  	}
   403  	if _, err = pngimg.WriteTo(f); err != nil {
   404  		log.Panic(err)
   405  	}
   406  }
   407
View as plain text