Support Sankey Diagrams (#69)

mfreeborn · web-flow · commit 2f7405fa016b · 2022-05-10T15:35:24.000+01:00
* working basic example

* fix examples

* implement the rest of sankey

* fix Domain

* improve test

* add sankey tests

* add test for Domain

* fix tests

* add Sankey to book
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Support for layout templates and pre-defined themes
 - Support for WASM environments
 - Lots of tests
+- Support for `Sankey` diagrams
 ### Changed
 - Improve implementation of `private::NumOrString` to support more primitive types ([Issue 
 #47](https://github.com/igiagkiozis/plotly/issues/47))
diff --git a/docs/book/src/recipes/basic_charts.md b/docs/book/src/recipes/basic_charts.md
@@ -7,3 +7,4 @@ Kind | Link
 Scatter Plots |[![Scatter Plots](./img/line_and_scatter_plot.png)](./basic_charts/scatter_plots.md)
 Line Charts | [![Line Charts](./img/line_shape_options_for_interpolation.png)](./basic_charts/line_charts.md)
 Bar Charts | [![Scatter Plots](./img/bar_chart_with_error_bars.png)](./basic_charts/scatter_plots.md)
+Sankey Diagrams | [![Scatter Plots](./img/basic_sankey.png)](./basic_charts/sankey_diagrams.md)
diff --git a/docs/book/src/recipes/basic_charts/sankey_diagrams.md b/docs/book/src/recipes/basic_charts/sankey_diagrams.md
@@ -0,0 +1,133 @@
+# Line Charts
+
+The following imports have been used to produce the plots below:
+
+```rust
+use itertools_num::linspace;
+use plotly::common::{
+    ColorScale, ColorScalePalette, DashType, Fill, Font, Line, LineShape, Marker, Mode, Title,
+};
+use plotly::layout::{Axis, BarMode, Layout, Legend, TicksDirection};
+use plotly::Sankey;
+use rand_distr::{Distribution, Normal, Uniform};
+```
+
+The `to_inline_html` method is used to produce the html plot displayed in this page.
+
+
+## Adding Names to Line and Scatter Plot
+```rust
+let trace = Sankey::new()
+        .orientation(Orientation::Horizontal)
+        .node(
+            Node::new()
+                .pad(15)
+                .thickness(30)
+                .line(SankeyLine::new().color(NamedColor::Black).width(0.5))
+                .label(vec!["A1", "A2", "B1", "B2", "C1", "C2"])
+                .color_array(vec![
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                ]),
+        )
+        .link(
+            Link::new()
+                .value(vec![8, 4, 2, 8, 4, 2])
+                .source(vec![0, 1, 0, 2, 3, 3])
+                .target(vec![2, 3, 3, 4, 4, 5]),
+        );
+
+    let layout = Layout::new()
+        .title("Basic Sankey".into())
+        .font(Font::new().size(10));
+
+    let mut plot = Plot::new();
+    plot.add_trace(trace);
+    plot.set_layout(layout);
+
+    if show {
+        plot.show();
+    }
+
+    println!("{}", plot.to_inline_html(Some("basic_sankey_diagram")));
+}
+```
+<div id="basic_sankey_diagram" class="plotly-graph-div" style="height:100%; width:100%;"></div>
+<script type="text/javascript">
+    window.PLOTLYENV=window.PLOTLYENV || {};
+    if (document.getElementById("basic_sankey_diagram")) {
+        var image_element = document.getElementById('image-export')
+
+        Plotly.newPlot('basic_sankey_diagram', {
+  "data": [
+    {
+      "type": "sankey",
+      "orientation": "h",
+      "node": {
+        "color": [
+          "blue",
+          "blue",
+          "blue",
+          "blue",
+          "blue",
+          "blue"
+        ],
+        "label": [
+          "A1",
+          "A2",
+          "B1",
+          "B2",
+          "C1",
+          "C2"
+        ],
+        "line": {
+          "color": "black",
+          "width": 0.5
+        },
+        "pad": 15,
+        "thickness": 30
+      },
+      "link": {
+        "source": [
+          0,
+          1,
+          0,
+          2,
+          3,
+          3
+        ],
+        "target": [
+          2,
+          3,
+          3,
+          4,
+          4,
+          5
+        ],
+        "value": [
+          8,
+          4,
+          2,
+          8,
+          4,
+          2
+        ]
+      }
+    }
+  ],
+  "layout": {
+    "title": {
+      "text": "Basic Sankey"
+    },
+    "font": {
+      "size": 10
+    }
+  },
+  "config": {}
+});
+    };
+</script>
diff --git a/docs/book/src/recipes/img/basic_sankey.png b/docs/book/src/recipes/img/basic_sankey.png
diff --git a/plotly/examples/basic_charts.rs b/plotly/examples/basic_charts.rs
@@ -1,9 +1,11 @@
 use itertools_num::linspace;
 use plotly::common::{
-    ColorScale, ColorScalePalette, DashType, Fill, Font, Line, LineShape, Marker, Mode, Title,
+    ColorScale, ColorScalePalette, DashType, Fill, Font, Line, LineShape, Marker, Mode,
+    Orientation, Title,
 };
 use plotly::layout::{Axis, BarMode, Layout, Legend, TicksDirection, TraceOrder};
-use plotly::{Bar, NamedColor, Plot, Rgb, Rgba, Scatter, ScatterPolar};
+use plotly::sankey::{Line as SankeyLine, Link, Node};
+use plotly::{Bar, NamedColor, Plot, Rgb, Rgba, Sankey, Scatter, ScatterPolar};
 use rand_distr::{Distribution, Normal, Uniform};
 
 // Scatter Plots
@@ -610,6 +612,48 @@ fn stacked_bar_chart(show: bool) {
     println!("{}", plot.to_inline_html(Some("stacked_bar_chart")));
 }
 
+// Sankey Diagrams
+fn basic_sankey_diagram(show: bool) {
+    // https://plotly.com/javascript/sankey-diagram/#basic-sankey-diagram
+    let trace = Sankey::new()
+        .orientation(Orientation::Horizontal)
+        .node(
+            Node::new()
+                .pad(15)
+                .thickness(30)
+                .line(SankeyLine::new().color(NamedColor::Black).width(0.5))
+                .label(vec!["A1", "A2", "B1", "B2", "C1", "C2"])
+                .color_array(vec![
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                    NamedColor::Blue,
+                ]),
+        )
+        .link(
+            Link::new()
+                .value(vec![8, 4, 2, 8, 4, 2])
+                .source(vec![0, 1, 0, 2, 3, 3])
+                .target(vec![2, 3, 3, 4, 4, 5]),
+        );
+
+    let layout = Layout::new()
+        .title("Basic Sankey".into())
+        .font(Font::new().size(10));
+
+    let mut plot = Plot::new();
+    plot.add_trace(trace);
+    plot.set_layout(layout);
+
+    if show {
+        plot.show();
+    }
+
+    println!("{}", plot.to_inline_html(Some("basic_sankey_diagram")));
+}
+
 fn main() -> std::io::Result<()> {
     // Scatter Plots
     simple_scatter_plot(true);
@@ -633,5 +677,8 @@ fn main() -> std::io::Result<()> {
     basic_bar_chart(true);
     grouped_bar_chart(true);
     stacked_bar_chart(true);
+
+    // Sankey Diagrams
+    basic_sankey_diagram(true);
     Ok(())
 }
diff --git a/plotly/src/common/mod.rs b/plotly/src/common/mod.rs
@@ -55,6 +55,65 @@ pub enum HoverInfo {
     Skip,
 }
 
+#[derive(Serialize, Clone, Debug, Default)]
+pub struct LegendGroupTitle {
+    text: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    font: Option<Font>,
+}
+
+impl LegendGroupTitle {
+    pub fn new(text: &str) -> Self {
+        Self {
+            text: text.to_string(),
+            ..Default::default()
+        }
+    }
+
+    pub fn font(mut self, font: Font) -> Self {
+        self.font = Some(font);
+        self
+    }
+}
+
+#[derive(Serialize, Clone, Debug, Default)]
+pub struct Domain {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    column: Option<usize>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    row: Option<usize>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    x: Option<[f64; 2]>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    y: Option<[f64; 2]>,
+}
+
+impl Domain {
+    pub fn new() -> Self {
+        Default::default()
+    }
+
+    pub fn column(mut self, column: usize) -> Self {
+        self.column = Some(column);
+        self
+    }
+
+    pub fn row(mut self, row: usize) -> Self {
+        self.row = Some(row);
+        self
+    }
+
+    pub fn x(mut self, x: &[f64; 2]) -> Self {
+        self.x = Some(x.to_owned());
+        self
+    }
+
+    pub fn y(mut self, y: &[f64; 2]) -> Self {
+        self.y = Some(y.to_owned());
+        self
+    }
+}
+
 #[derive(Serialize, Clone, Debug)]
 #[serde(rename_all = "lowercase")]
 pub enum TextPosition {
@@ -74,9 +133,10 @@ pub enum ConstrainText {
 }
 
 #[derive(Serialize, Clone, Debug)]
-#[serde(rename_all = "lowercase")]
 pub enum Orientation {
+    #[serde(rename = "v")]
     Vertical,
+    #[serde(rename = "h")]
     Horizontal,
 }
 
@@ -148,6 +208,7 @@ pub enum PlotType {
     Histogram,
     Histogram2dContour,
     Ohlc,
+    Sankey,
     Surface,
 }
 
@@ -1472,6 +1533,19 @@ mod tests {
     use super::*;
     use crate::NamedColor;
 
+    #[test]
+    fn test_serialize_domain() {
+        let domain = Domain::new().column(0).row(0).x(&[0., 1.]).y(&[0., 1.]);
+        let expected = json!({
+            "column": 0,
+            "row": 0,
+            "x": [0.0, 1.0],
+            "y": [0.0, 1.0],
+        });
+
+        assert_eq!(to_value(domain).unwrap(), expected);
+    }
+
     #[test]
     fn test_serialize_direction() {
         // TODO: I think `Direction` would be better as a struct, with `fillcolor` and `line` attributes
@@ -1519,8 +1593,8 @@ mod tests {
     #[test]
     #[rustfmt::skip]
     fn test_serialize_orientation() {
-        assert_eq!(to_value(Orientation::Vertical).unwrap(), json!("vertical"));
-        assert_eq!(to_value(Orientation::Horizontal).unwrap(), json!("horizontal"));
+        assert_eq!(to_value(Orientation::Vertical).unwrap(), json!("v"));
+        assert_eq!(to_value(Orientation::Horizontal).unwrap(), json!("h"));
     }
 
     #[test]
@@ -1583,6 +1657,7 @@ mod tests {
         assert_eq!(to_value(PlotType::Histogram).unwrap(), json!("histogram"));
         assert_eq!(to_value(PlotType::Histogram2dContour).unwrap(), json!("histogram2dcontour"));
         assert_eq!(to_value(PlotType::Ohlc).unwrap(), json!("ohlc"));
+        assert_eq!(to_value(PlotType::Sankey).unwrap(), json!("sankey"));
         assert_eq!(to_value(PlotType::Surface).unwrap(), json!("surface"));
     }
 
diff --git a/plotly/src/layout.rs b/plotly/src/layout.rs
@@ -3513,7 +3513,7 @@ mod tests {
             "bordercolor": "#321321",
             "borderwidth": 500,
             "font": {},
-            "orientation": "vertical",
+            "orientation": "v",
             "traceorder": "normal",
             "tracegroupgap": 10,
             "itemsizing": "trace",
@@ -3980,7 +3980,7 @@ mod tests {
             .color("#000FFF")
             .active_color("#AAABBB");
         let expected = json!({
-            "orientation": "horizontal",
+            "orientation": "h",
             "bgcolor": "#FFF000",
             "color": "#000FFF",
             "activecolor": "#AAABBB",
diff --git a/plotly/src/lib.rs b/plotly/src/lib.rs
@@ -26,6 +26,7 @@ pub mod contour;
 pub mod heat_map;
 pub mod histogram;
 pub mod ohlc;
+pub mod sankey;
 pub mod scatter;
 pub mod scatter_polar;
 pub mod surface;
@@ -42,6 +43,7 @@ pub use crate::contour::Contour;
 pub use crate::heat_map::HeatMap;
 pub use crate::histogram::Histogram;
 pub use crate::ohlc::Ohlc;
+pub use crate::sankey::Sankey;
 pub use crate::scatter::Scatter;
 pub use crate::scatter_polar::ScatterPolar;
 pub use crate::surface::Surface;
diff --git a/plotly/src/sankey.rs b/plotly/src/sankey.rs
