Merge pull request #41 from huijunchen9260/master

add wireframe based on plot_surface

Author: jacobwilliams

src/pyplot_module.F90

@@ -24,7 +24,7 @@
 
     module pyplot_module
 
-    use, intrinsic :: iso_fortran_env 
+    use, intrinsic :: iso_fortran_env
 
     implicit none
 
@@ -83,6 +83,7 @@ module pyplot_module
         procedure, public :: add_3d_plot   !! add a 3d plot to pyplot instance
         procedure, public :: add_sphere    !! add a 3d sphere to pyplot instance
         procedure, public :: add_contour   !! add a contour plot to pyplot instance
+        procedure, public :: plot_wireframe!! add a wireframe plot to pyplot instance
         procedure, public :: plot_surface  !! add a surface plot to pyplot instance
         procedure, public :: add_bar       !! add a barplot to pyplot instance
         procedure, public :: add_imshow    !! add an image plot (using `imshow`)
@@ -656,6 +657,96 @@ subroutine plot_surface(me, x, y, z, label, linestyle, linewidth, levels, color,
     end subroutine plot_surface
 !*****************************************************************************************
 
+!*****************************************************************************************
+!> author: Jacob Williams
+!
+! Add a wireframe plot.
+!
+!@note This requires `use_numpy` to be True.
+
+    subroutine plot_wireframe(me, x, y, z, label, linestyle, linewidth, levels, color, &
+                            cmap, colorbar, antialiased, istat)
+
+    class(pyplot),           intent (inout)         :: me           !! pyplot handler
+    real(wp),dimension(:),   intent (in)            :: x            !! x values
+    real(wp),dimension(:),   intent (in)            :: y            !! y values
+    real(wp),dimension(:,:), intent (in)            :: z            !! z values (a matrix)
+    character(len=*),        intent (in)            :: label        !! plot label
+    character(len=*),        intent (in)            :: linestyle    !! style of the plot line
+    integer,                 intent (in),  optional :: linewidth    !! width of the plot line
+    real(wp),dimension(:),   intent (in),  optional :: levels       !! contour levels to plot
+    character(len=*),        intent (in),  optional :: color        !! Color of the surface patches
+    character(len=*),        intent (in),  optional :: cmap         !! colormap if filled=True (examples: 'jet', 'bone')
+    logical,                 intent (in),  optional :: colorbar     !! add a colorbar (default=False)
+    logical,                 intent (in),  optional :: antialiased  !! The surface is made opaque by using antialiased=False
+    integer,                 intent (out), optional :: istat        !! status output (0 means no problems)
+
+    character(len=:), allocatable :: xstr           !! x values stringified
+    character(len=:), allocatable :: ystr           !! y values stringified
+    character(len=:), allocatable :: zstr           !! z values stringified
+    character(len=:), allocatable :: levelstr       !! levels vector stringified
+    character(len=:), allocatable :: antialiasedstr !! antialiased stringified
+    character(len=max_int_len)    :: iline          !! actual line width
+    character(len=*), parameter   :: xname = 'x'    !! x variable name for script
+    character(len=*), parameter   :: yname = 'y'    !! y variable name for script
+    character(len=*), parameter   :: zname = 'z'    !! z variable name for script
+    character(len=*), parameter   :: xname_ = 'X'   !! X variable name for contour
+    character(len=*), parameter   :: yname_ = 'Y'   !! Y variable name for contour
+    character(len=*), parameter   :: zname_ = 'Z'   !! Z variable name for contour
+    character(len=:), allocatable :: extras         !! optional stuff
+
+    if (allocated(me%str)) then
+
+        if (present(istat)) istat = 0
+
+        !convert the arrays to strings:
+        call vec_to_string(x, me%real_fmt, xstr, me%use_numpy)
+        call vec_to_string(y, me%real_fmt, ystr, me%use_numpy)
+        call matrix_to_string(z, me%real_fmt, zstr, me%use_numpy)
+        if (present(levels)) call vec_to_string(levels, me%real_fmt, levelstr, me%use_numpy)
+
+        !get optional inputs (if not present, set default value):
+        call optional_int_to_string(linewidth, iline, '3')
+        call optional_logical_to_string(antialiased, antialiasedstr, 'False')
+
+        !write the arrays:
+        call me%add_str(trim(xname)//' = '//xstr)
+        call me%add_str(trim(yname)//' = '//ystr)
+        call me%add_str(trim(zname)//' = '//zstr)
+        call me%add_str('')
+
+        !convert inputs for contour plotting:
+        call me%add_str(xname_//', '//yname_//' = np.meshgrid('//trim(xname)//', '//trim(yname)//')')
+        call me%add_str(zname_//' = np.transpose('//zname//')')
+
+        !optional arguments:
+        extras = ''
+        if (present(levels))     extras = extras//','//'levels='//levelstr
+        if (present(color))      extras = extras//','//'colors='//trim(me%raw_str_token)//'"'//color//'"'
+        if (present(linewidth))  extras = extras//','//'linewidths='//trim(adjustl(iline))
+        if (present(cmap))       extras = extras//','//'cmap='//trim(me%raw_str_token)//'"'//cmap//'"'
+
+        !write the plot statement:
+        call me%add_str('CS = ax.plot_wireframe'//'('//xname_//','//yname_//','//zname_//','//&
+                        'label='//trim(me%raw_str_token)//'"'//trim(label)//'",'//&
+                        'antialiased='//trim(antialiasedstr)//','//&
+                        'linestyles='//trim(me%raw_str_token)//'"'//trim(adjustl(linestyle))//'"'//&
+                        extras//')')
+
+        if (present(colorbar)) then
+            if (colorbar) call me%add_str('fig.colorbar(CS)')
+        end if
+
+        call me%add_str('')
+
+    else
+        if (present(istat)) istat = -1
+        write(error_unit,'(A)') 'Error in add_plot: pyplot class not properly initialized.'
+    end if
+
+    end subroutine plot_wireframe
+!*****************************************************************************************
+
 !*****************************************************************************************
 !> author: Jacob Williams
 !