Improve setting of axis labels. Impl set axis labels 3d

src/constants.rs

@@ -23,6 +23,7 @@
 ///   For example a circle will show as a circle in the screen and not an ellipse. This function also handles
 ///   the 3D case which is a little tricky with Matplotlib. In this case (3D), the version of Matplotlib
 ///   must be greater than 3.3.0.
+/// * `set_axis_label` -- Sets the label of the axis along the dimension 'dim'
 pub const PYTHON_HEADER: &str = "### file generated by the 'plotpy' Rust crate
 
 import numpy as np
@@ -68,6 +69,7 @@ def ax3d():
 def subplot3d(a,b,c):
     global THREE_D_ACTIVE
     THREE_D_ACTIVE = (a,b,c)
+    ax3d()
 
 # Transforms data limits to axis limits
 def data_to_axis(coords):
@@ -101,6 +103,17 @@ def set_equal_axes():
         print('VERSION of MATPLOTLIB = {}'.format(matplotlib.__version__))
         print('ERROR: set_box_aspect is missing in this version of Matplotlib')
 
+# Sets the label of the axis along the dimension 'dim'
+def set_axis_label(dim, label):
+    global THREE_D
+    if len(THREE_D) == 0:
+        if dim == 1: plt.gca().set_xlabel(label)
+        if dim == 2: plt.gca().set_ylabel(label)
+    else:
+        if dim == 1: ax3d().set_xlabel(label)
+        if dim == 2: ax3d().set_ylabel(label)
+        if dim == 3: ax3d().set_zlabel(label)
+
 ################## plotting commands follow after this line ############################
 
 ";
@@ -144,6 +157,6 @@ mod tests {
 
     #[test]
     fn constants_are_correct() {
-        assert_eq!(PYTHON_HEADER.len(), 2729);
+        assert_eq!(PYTHON_HEADER.len(), 3119);
     }
 }

src/plot.rs

@@ -160,8 +160,8 @@ impl Plot {
             &mut self.buffer,
             "plt.gca().set_axisbelow(True)\n\
              plt.grid(linestyle='--',color='grey',zorder=-1000)\n\
-             plt.xlabel(r'{}')\n\
-             plt.ylabel(r'{}')\n",
+             set_axis_label(1,r'{}')\n\
+             set_axis_label(2,r'{}')\n",
             xlabel, ylabel
         )
         .unwrap();
@@ -174,8 +174,8 @@ impl Plot {
             &mut self.buffer,
             "plt.gca().set_axisbelow(True)\n\
              plt.grid(linestyle='--',color='grey',zorder=-1000)\n\
-             plt.xlabel(r'{}')\n\
-             plt.ylabel(r'{}')\n",
+             set_axis_label(1,r'{}')\n\
+             set_axis_label(2,r'{}')\n",
             xlabel, ylabel
         )
         .unwrap();
@@ -635,27 +635,44 @@ impl Plot {
 
     /// Sets the label for the x-axis
     pub fn set_label_x(&mut self, label: &str) -> &mut Self {
-        write!(&mut self.buffer, "plt.xlabel(r'{}')\n", label).unwrap();
+        write!(&mut self.buffer, "set_axis_label(1,r'{}')\n", label).unwrap();
         self
     }
 
     /// Sets the label for the y-axis
     pub fn set_label_y(&mut self, label: &str) -> &mut Self {
-        write!(&mut self.buffer, "plt.ylabel(r'{}')\n", label).unwrap();
+        write!(&mut self.buffer, "set_axis_label(2,r'{}')\n", label).unwrap();
         self
     }
 
-    /// Sets the labels of x and y axis
+    /// Sets the label for the z-axis
+    pub fn set_label_z(&mut self, label: &str) -> &mut Self {
+        write!(&mut self.buffer, "set_axis_label(3,r'{}')\n", label).unwrap();
+        self
+    }
+
+    /// Sets the labels for the x and y axes
     pub fn set_labels(&mut self, xlabel: &str, ylabel: &str) -> &mut Self {
         write!(
             &mut self.buffer,
-            "plt.xlabel(r'{}')\nplt.ylabel(r'{}')\n",
+            "set_axis_label(1,r'{}')\nset_axis_label(2,r'{}')\n",
             xlabel, ylabel
         )
         .unwrap();
         self
     }
 
+    /// Sets the labels for the x, y, and z axes
+    pub fn set_labels_3d(&mut self, xlabel: &str, ylabel: &str, zlabel: &str) -> &mut Self {
+        write!(
+            &mut self.buffer,
+            "set_axis_label(1,r'{}')\nset_axis_label(2,r'{}')\nset_axis_label(2,r'{}')\n",
+            xlabel, ylabel, zlabel
+        )
+        .unwrap();
+        self
+    }
+
     /// Sets inverted x-axis
     pub fn set_inv_x(&mut self) -> &mut Self {
         write!(&mut self.buffer, "plt.gca().invert_xaxis()\n").unwrap();
@@ -858,12 +875,12 @@ mod tests {
         plot.grid_and_labels("xx", "yy").grid_labels_legend("xx", "yy").legend();
         let b: &str = "plt.gca().set_axisbelow(True)\n\
                        plt.grid(linestyle='--',color='grey',zorder=-1000)\n\
-                       plt.xlabel(r'xx')\n\
-                       plt.ylabel(r'yy')\n\
+                       set_axis_label(1,r'xx')\n\
+                       set_axis_label(2,r'yy')\n\
                        plt.gca().set_axisbelow(True)\n\
                        plt.grid(linestyle='--',color='grey',zorder=-1000)\n\
-                       plt.xlabel(r'xx')\n\
-                       plt.ylabel(r'yy')\n\
+                       set_axis_label(1,r'xx')\n\
+                       set_axis_label(2,r'yy')\n\
                        h,l=plt.gca().get_legend_handles_labels()\n\
                        if len(h)>0 and len(l)>0:\n\
                        \x20\x20\x20\x20leg=plt.legend(handlelength=3,ncol=1,loc='best')\n\
@@ -933,10 +950,10 @@ mod tests {
                        plt.gca().set_yscale('log')\n\
                        plt.gca().set_xscale('linear')\n\
                        plt.gca().set_yscale('linear')\n\
-                       plt.xlabel(r'x-label')\n\
-                       plt.ylabel(r'y-label')\n\
-                       plt.xlabel(r'x')\n\
-                       plt.ylabel(r'y')\n\
+                       set_axis_label(1,r'x-label')\n\
+                       set_axis_label(2,r'y-label')\n\
+                       set_axis_label(1,r'x')\n\
+                       set_axis_label(2,r'y')\n\
                        plt.gca().view_init(elev=1,azim=10)\n\
                        major_locator = tck.MultipleLocator(1.5)\n\
                        n_ticks = (plt.gca().axis()[1] - plt.gca().axis()[0]) / 1.5\n\
@@ -1119,4 +1136,15 @@ mod tests {
                        plt.gcf().align_labels()\n";
         assert_eq!(plot.buffer, b);
     }
+
+    #[test]
+    fn set_labels_3d_works() {
+        let mut plot = Plot::new();
+        plot.set_label_z("Z").set_labels_3d("X", "Y", "Z");
+        let b: &str = "set_axis_label(3,r'Z')\n\
+                       set_axis_label(1,r'X')\n\
+                       set_axis_label(2,r'Y')\n\
+                       set_axis_label(2,r'Z')\n";
+        assert_eq!(plot.buffer, b);
+    }
 }

tests/test_subplot3d.rs

@@ -34,10 +34,24 @@ fn test_subplot3d() -> Result<(), StrError> {
         .draw(&x, &y, &z);
 
     // add surfaces to plot
-    plot.set_subplot3d(2, 2, 1).add(&surf1);
-    plot.set_subplot3d(2, 2, 2).add(&surf2);
-    plot.set_subplot3d(2, 2, 3).add(&surf3);
-    plot.set_subplot3d(2, 2, 4).add(&surf4);
+    plot.set_subplot3d(2, 2, 1)
+        .set_label_x("X AXIS IS BEAUTIFUL")
+        .set_label_y("Y AXIS IS BEAUTIFUL")
+        .set_label_z("Z AXIS IS BEAUTIFUL")
+        .add(&surf1);
+    plot.set_subplot3d(2, 2, 2)
+        .set_label_x("X AXIS IS BEAUTIFUL")
+        .set_label_y("Y AXIS IS BEAUTIFUL")
+        .set_label_z("Z AXIS")
+        .add(&surf2);
+    plot.set_subplot3d(2, 2, 3)
+        .set_label_x("X AXIS IS BEAUTIFUL")
+        .set_label_y("Y AXIS IS BEAUTIFUL")
+        .set_label_z("Z AXIS")
+        .add(&surf3);
+    plot.set_subplot3d(2, 2, 4)
+        .set_labels_3d("X AXIS IS BEAUTIFUL", "Y AXIS IS BEAUTIFUL", "Z AXIS IS BEAUTIFUL")
+        .add(&surf4);
 
     // save figure
     let path = Path::new(OUT_DIR).join("integ_subplot3d.svg");
@@ -47,7 +61,7 @@ fn test_subplot3d() -> Result<(), StrError> {
     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 > 2900 && n < 3100);
+    let n_lines = lines_iter.count();
+    assert!(n_lines > 3300 && n < 3400);
     Ok(())
 }