diff --git a/examples/README.md b/examples/README.md
index 9d4f905..1cac750 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -64,6 +64,7 @@ Some output of integration tests are shown below.
[test_contour.rs](https://github.com/cpmech/plotpy/tree/main/tests/test_contour.rs)
+
diff --git a/figures/integ_contour_colorbar_options.svg b/figures/integ_contour_colorbar_options.svg
new file mode 100644
index 0000000..2e8f725
--- /dev/null
+++ b/figures/integ_contour_colorbar_options.svg
@@ -0,0 +1,2188 @@
+
+
+
diff --git a/src/constants.rs b/src/constants.rs
index 35754d4..5d38e4d 100644
--- a/src/constants.rs
+++ b/src/constants.rs
@@ -35,6 +35,7 @@ import matplotlib.lines as lns
import matplotlib.transforms as tra
import mpl_toolkits.mplot3d
import matplotlib.tri as plt_tri
+from mpl_toolkits.axes_grid1 import make_axes_locatable
# Variable to handle NaN values coming from Rust
NaN = np.nan
@@ -151,6 +152,6 @@ mod tests {
#[test]
fn constants_are_correct() {
- assert_eq!(PYTHON_HEADER.len(), 2919);
+ assert_eq!(PYTHON_HEADER.len(), 2975);
}
}
diff --git a/src/contour.rs b/src/contour.rs
index 35e90c8..f0afabf 100644
--- a/src/contour.rs
+++ b/src/contour.rs
@@ -57,6 +57,9 @@ pub struct Contour {
no_inline_labels: bool, // Do not draw labels inline
no_colorbar: bool, // Skip drawing a colorbar
colorbar_label: String, // Colorbar label
+ colorbar_axes: String, // Axes into which the colorbar will be drawn
+ colorbar_location: String, // Colorbar location
+ colorbar_extra: String, // Extra options for the colorbar
number_format_cb: String, // Number format for the labels in lines contour
line_color: String, // Line color for the lines contour
line_style: String, // Line style for the lines contour
@@ -89,6 +92,9 @@ impl Contour {
no_inline_labels: false,
no_colorbar: false,
colorbar_label: String::new(),
+ colorbar_axes: String::new(),
+ colorbar_location: String::new(),
+ colorbar_extra: String::new(),
number_format_cb: String::new(),
line_color: "black".to_string(),
line_style: String::new(),
@@ -201,7 +207,12 @@ impl Contour {
}
if !self.no_colorbar && !self.no_fill {
let opt_colorbar = self.options_colorbar();
- write!(&mut self.buffer, "cb=plt.colorbar(cf{})\n", &opt_colorbar).unwrap();
+ if self.colorbar_axes != "" {
+ write!(&mut self.buffer, "{}", self.colorbar_axes).unwrap();
+ write!(&mut self.buffer, "cb=plt.colorbar(cf,cax=cax{})\n", &opt_colorbar).unwrap();
+ } else {
+ write!(&mut self.buffer, "cb=plt.colorbar(cf{})\n", &opt_colorbar).unwrap();
+ }
if self.colorbar_label != "" {
write!(&mut self.buffer, "cb.ax.set_ylabel(r'{}')\n", self.colorbar_label).unwrap();
}
@@ -294,6 +305,42 @@ impl Contour {
self
}
+ /// Configure the axes into which the colorbar will be drawn
+ ///
+ /// # Arguments
+ ///
+ /// * `location` -- location of the colorbar axes (e.g., 'right', 'top', 'left', 'bottom')
+ /// * `width_pct` -- width percentage of the colorbar axes (e.g., 5.0, 10.0, 15.0)
+ /// * `pad` -- padding between the main plot and the colorbar axes (e.g., 0.1, 0.2)
+ pub fn set_colorbar_axes(&mut self, location: &str, width_pct: f64, pad: f64) -> &mut Self {
+ self.colorbar_axes = format!(
+ "ax=plt.gca()\n\
+ divider=make_axes_locatable(ax)\n\
+ cax=divider.append_axes('{}',size='{}%',pad={})\n\
+ plt.sca(ax)\n",
+ location, width_pct, pad
+ );
+ self
+ }
+
+ /// Sets the colorbar location
+ ///
+ /// Options: 'right', 'left', 'top', 'bottom'
+ ///
+ /// See:
+ pub fn set_colorbar_location(&mut self, location: &str) -> &mut Self {
+ self.colorbar_location = location.to_string();
+ self
+ }
+
+ /// Sets extra options for the colorbar
+ ///
+ /// Example `extra = "fraction=0.046, pad=0.04"`
+ pub fn set_colorbar_extra(&mut self, extra: &str) -> &mut Self {
+ self.colorbar_extra = String::from(extra);
+ self
+ }
+
/// Sets the number format for the labels in the colorbar (cb)
pub fn set_number_format_cb(&mut self, format: &str) -> &mut Self {
self.number_format_cb = String::from(format);
@@ -471,6 +518,12 @@ impl Contour {
if self.number_format_cb != "" {
write!(&mut opt, ",format='{}'", self.number_format_cb).unwrap();
}
+ if self.colorbar_location != "" {
+ write!(&mut opt, ",location='{}'", self.colorbar_location).unwrap();
+ }
+ if self.colorbar_extra != "" {
+ write!(&mut opt, ",{}", self.colorbar_extra).unwrap();
+ }
opt
}
diff --git a/tests/test_contour.rs b/tests/test_contour.rs
index 1e099c8..1040ea1 100644
--- a/tests/test_contour.rs
+++ b/tests/test_contour.rs
@@ -150,7 +150,6 @@ fn test_contour_colormap_name() -> Result<(), StrError> {
#[test]
fn test_contour_draw_tri() -> Result<(), StrError> {
// point coordinates (two triangles in a square)
-
let x = vec![0.0, 1.0, 2.0, 3.0, 0.5, 1.5, 2.5, 1.0, 2.0, 1.5];
let y = vec![0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 2.0, 2.0, 3.0];
let connectivity = vec![
@@ -170,7 +169,7 @@ fn test_contour_draw_tri() -> Result<(), StrError> {
z[i] = x[i] + y[i];
}
- // canvas
+ // contour
let mut contour = Contour::new();
contour
.set_no_fill(false)
@@ -203,3 +202,47 @@ fn test_contour_draw_tri() -> Result<(), StrError> {
assert!(n > 1030 && n < 1090);
Ok(())
}
+
+#[test]
+fn test_contour_colorbar_options() -> Result<(), StrError> {
+ // data
+ let n = 9;
+ let (x, y, z) = generate3d(-2.0, 2.0, 0.0, 2.0, n, n, |x, y| x * x + y * y);
+
+ // contour A
+ let mut contour_a = Contour::new();
+ contour_a
+ .set_colorbar_axes("right", 5.0, 0.1)
+ .set_colorbar_label("x² + y²")
+ .draw(&x, &y, &z);
+
+ // contour B
+ let mut contour_b = Contour::new();
+ contour_b
+ .set_colorbar_extra("fraction=0.12,pad=0.12")
+ .set_colorbar_location("bottom")
+ .set_colorbar_label("x² + y²")
+ .set_colors(&vec!["#fcaeae", "#da98d1", "#c45178", "#5594d2", "#e6af69", "#e6d969"])
+ .draw(&x, &y, &z);
+
+ // add contour to plot
+ let mut plot = Plot::new();
+ plot.set_gaps(0.3, 0.0);
+ plot.set_subplot(1, 2, 1).set_equal_axes(true).add(&contour_a);
+ plot.set_subplot(1, 2, 2).set_equal_axes(true).add(&contour_b);
+
+ // save figure
+ let path = Path::new(OUT_DIR).join("integ_contour_colorbar_options.svg");
+ plot.set_figure_size_points(600.0, 300.0)
+ .set_equal_axes(true)
+ .set_show_errors(true)
+ .save(&path)?;
+
+ // check number of lines
+ let file = File::open(path).map_err(|_| "cannot open file")?;
+ let buffered = BufReader::new(file);
+ let lines_iter = buffered.lines();
+ let n = lines_iter.count();
+ assert!(n > 2150 && n < 2200);
+ Ok(())
+}