Add bar plots

Author: Chris00

examples/plot_types.rs

@@ -1,6 +1,7 @@
 /// Inspired from
 /// https://matplotlib.org/stable/plot_types/index.html#plot-types
-use matplotlib::{self as mpl, figure::Figure};
+use matplotlib::{colors::Base, self as mpl, figure::Figure};
+use ndarray::Array1;
 
 /// Base name for all the PDF files generated by this script.
 static BASE: &'static str = "target/plot_types_";
@@ -9,6 +10,7 @@ fn main() -> anyhow::Result<()> {
     // Pairwise data
     plot_xy()?;
     scatter_xy()?;
+    bar()?;
     Ok(())
 }
 
@@ -62,3 +64,20 @@ fn scatter_xy() -> anyhow::Result<()> {
     fig.save().to_file(format!("{BASE}scatter_xy.pdf"))?;
     Ok(())
 }
+
+fn bar() -> anyhow::Result<()> {
+    mpl::style::using("_mpl-gallery")?;
+
+    let x: Vec<_> = (0 .. 8).map(|x| 0.5 + x as f64).collect();
+    let y = [4.8, 5.5, 3.5, 4.6, 6.5, 6.6, 2.6, 3.0];
+
+    let fig = Figure::new()?;
+    let [[mut ax]] = fig.subplots()?;
+
+    ax.bar(&x, &y).width(1.).edgecolor(Base::W).linewidth(0.7).plot();
+
+    ax.set_xlim(0., 8.)   .set_xticks(Array1::range(1., 8., 1.))
+        .set_ylim(0., 8.) .set_yticks(Array1::range(1., 8., 1.));
+    fig.save().to_file(format!("{BASE}bar.pdf"))?;
+    Ok(())
+}

src/axes.rs

@@ -211,11 +211,24 @@ impl Axes {
 
     /// Scatter plot of `y` vs. `x` with optional varying marker size
     /// and/or color.
-    pub fn scatter<'a, D>(&'a mut self, x: D, y: D) -> Scatter<'a, D>
+    pub fn scatter<'a, D>(&'a mut self, x: D, height: D) -> Scatter<'a, D>
     where
         D: AsRef<[f64]>,
     {
-        Scatter::new(self, x, y)
+        Scatter::new(self, x, height)
+    }
+
+    /// Make a bar plot.
+    ///
+    /// The bars are positioned at `x` with the given
+    /// [alignment][`Bar::align`].  Their dimensions are given by
+    /// height and width. The vertical baseline is bottom (default 0).
+    pub fn bar<'a, D1, D2>(&'a mut self, x: D1, height: D2) -> Bar<'a>
+    where
+        D1: AsRef<[f64]> + 'a,
+        D2: AsRef<[f64]> + 'a,
+    {
+        Bar::new(self, x, height)
     }
 
     /// Set the title to `txt` for the Axes.
@@ -661,6 +674,7 @@ where
     }
 }
 
+/// Options for [`Axes::scatter`].
 #[must_use]
 pub struct Scatter<'a, D> {
     axes: &'a Axes,
@@ -702,7 +716,7 @@ where D: AsRef<[f64]> {
         self
     }
 
-    /// Specify the marker colors.
+    /// Specify the marker color(s).
     pub fn c<C>(mut self, colors: impl ScatterColors) -> Self
     where C: Color,
     {
@@ -827,6 +841,131 @@ impl<C: Color> ScatterColors for C {
     }
 }
 
+pub struct Bar<'a> {
+    axes: &'a Axes,
+    x: Box<dyn AsRef<[f64]> + 'a>, // FIXME: categorical data ?
+    height: Box<dyn AsRef<[f64]> + 'a>,
+    width: f64, // FIXME: or array
+    bottom: f64, // FIXME: or array
+    align: BarAlign,
+    // Options
+    color: Option<[f64; 4]>, // FIXME: or array
+    facecolor: Option<[f64; 4]>, // FIXME: or array
+    edgecolor: Option<[f64; 4]>, // FIXME: or array
+    linewidth: Option<f64>, // FIXME: or array
+    tick_label: Option<&'a str>, // FIXME: or array
+    label: Option<&'a str>, // FIXME: or array
+    //xerr,  yerr, ecolor, capsize, error_kw, log
+}
+
+/// Alignment of [`Axes::bar`].  See [`Bar::align`].
+#[derive(Debug, Clone, Copy)]
+pub enum BarAlign {
+    Center,
+    Edge,
+}
+
+impl<'a> Bar<'a> {
+    fn new<D1, D2>(axes: &'a Axes, x: D1, height: D2) -> Self
+    where
+        D1: AsRef<[f64]> + 'a,
+        D2: AsRef<[f64]> + 'a,
+    {
+        Self {
+            axes,
+            x: Box::new(x),
+            height: Box::new(height),
+            width: 0.8,
+            bottom: 0.,
+            align: BarAlign::Center,
+            color: None,  facecolor: None,  edgecolor: None,
+            linewidth: None,  tick_label: None,  label: None,
+        }
+    }
+
+    pub fn width(mut self, w: f64) -> Self {
+        self.width = w;
+        self
+    }
+
+    pub fn bottom(mut self, b: f64) -> Self {
+        self.bottom = b;
+        self
+    }
+
+    pub fn align(mut self, a: BarAlign) -> Self {
+        self.align = a;
+        self
+    }
+
+    pub fn color(mut self, c: impl Color) -> Self {
+        self.color = Some(c.rgba());
+        self
+    }
+
+    pub fn facecolor(mut self, c: impl Color) -> Self {
+        self.facecolor = Some(c.rgba());
+        self
+    }
+
+    pub fn edgecolor(mut self, c: impl Color) -> Self {
+        self.edgecolor = Some(c.rgba());
+        self
+    }
+
+    pub fn linewidth(mut self, w: f64) -> Self {
+        self.linewidth = Some(w);
+        self
+    }
+
+    pub fn tick_label(mut self, l: &'a str) -> Self {
+        self.tick_label = Some(l);
+        self
+    }
+
+    pub fn label(mut self, l: &'a str) -> Self {
+        self.label = Some(l);
+        self
+    }
+
+    pub fn plot(self) {
+        Python::attach(|py| {
+            let x = self.x.as_ref().as_ref().to_pyarray(py);
+            let height = self.height.as_ref().as_ref().to_pyarray(py);
+            let align = match self.align {
+                BarAlign::Center => "center",
+                BarAlign::Edge => "edge",
+            };
+            let kwargs = PyDict::new(py);
+            kwargs.set_item("align", align).unwrap();
+            if let Some(color) = self.color {
+                kwargs.set_item("color", color).unwrap();
+            }
+            if let Some(facecolor) = self.facecolor {
+                kwargs.set_item("facecolor", facecolor).unwrap();
+            }
+            if let Some(edgecolor) = self.edgecolor {
+                kwargs.set_item("edgecolor", edgecolor).unwrap();
+            }
+            if let Some(lw) = self.linewidth {
+                kwargs.set_item("linewidth", lw).unwrap();
+            }
+            if let Some(tick_label) = self.tick_label {
+                kwargs.set_item("tick_label", tick_label).unwrap();
+            }
+            if let Some(label) = self.label {
+                kwargs.set_item("label", label).unwrap();
+            }
+            // TODO: options
+            self.axes.ax.bind(py)
+                .call_method(intern!(py, "bar"),
+                    (x, height, self.width, self.bottom),
+                    Some(&kwargs))
+                .unwrap();
+        })
+    }
+}
+
 pub struct QuadContourSet {
     contours: Py<PyAny>,
 }