Merge pull request #17 from jacobwilliams/sphere

Sphere

Author: jacobwilliams

src/pyplot_module.f90

@@ -51,6 +51,7 @@ module pyplot_module
 
         procedure, public :: add_plot      !! add a 2d plot to pyplot instance
         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 :: add_bar       !! add a barplot to pyplot instance
         procedure, public :: add_imshow    !! add an image plot (using `imshow`)
@@ -549,6 +550,83 @@ subroutine add_3d_plot(me, x, y, z, label, linestyle, markersize, linewidth, ist
     end subroutine add_3d_plot
 !*****************************************************************************************
 
+!*****************************************************************************************
+!> author: Jacob Williams
+!
+! Add a sphere to a 3D x,y,z plot.
+!
+!@note Must initialize the class with `mplot3d=.true.` and `use_numpy=.true.`.
+
+    subroutine add_sphere(me, r, xc, yc, zc, n_facets, linewidth, antialiased, color, istat)
+
+    implicit none
+
+    class(pyplot),    intent (inout)        :: me            !! pyplot handler
+    real(wp),         intent (in)           :: r             !! radius of the sphere
+    real(wp),         intent (in)           :: xc            !! x value of sphere center
+    real(wp),         intent (in)           :: yc            !! y value of sphere center
+    real(wp),         intent (in)           :: zc            !! z value of sphere center
+    integer,          intent (in), optional :: n_facets      !! [default is 100]
+    integer,          intent (in), optional :: linewidth     !! line width
+    logical,          intent (in), optional :: antialiased   !! enabled anti-aliasing
+    character(len=*), intent (in), optional :: color         !! color of the contour line
+    integer,          intent (out)          :: istat         !! status output (0 means no problems)
+
+    character(len=:), allocatable :: rstr         !! `r` value stringified
+    character(len=:), allocatable :: xcstr        !! `xc` value stringified
+    character(len=:), allocatable :: ycstr        !! `yc` value stringified
+    character(len=:), allocatable :: zcstr        !! `zc` value stringified
+    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=max_int_len)    :: linewidth_str       !! `linewidth` input stringified
+    character(len=:), allocatable :: antialiased_str     !! `antialised` input stringified
+    character(len=max_int_len)    :: n_facets_str        !! `n_facets` input stringified
+    character(len=:), allocatable :: extras              !! optional stuff string
+
+    if (allocated(me%str)) then
+
+        !get optional inputs (if not present, set default value):
+        call optional_int_to_string(n_facets, n_facets_str, '100')
+        extras = ''
+        if (present(linewidth)) then
+            call optional_int_to_string(linewidth, linewidth_str, '1')
+            extras = extras//','//'linewidth='//linewidth_str
+        end if
+        if (present(antialiased)) then
+            call optional_logical_to_string(antialiased, antialiased_str, 'False')
+            extras = extras//','//'antialiased='//antialiased_str
+        end if
+        if (present(color)) then
+            extras = extras//','//'color="'//trim(color)//'"'
+        end if
+
+        istat = 0
+
+        !convert the arrays to strings:
+        call real_to_string(r , me%real_fmt, rstr)
+        call real_to_string(xc, me%real_fmt, xcstr)
+        call real_to_string(yc, me%real_fmt, ycstr)
+        call real_to_string(zc, me%real_fmt, zcstr)
+
+        ! sphere code:
+        call me%add_str('u = np.linspace(0, 2 * np.pi, '//n_facets_str//')')
+        call me%add_str('v = np.linspace(0, np.pi, '//n_facets_str//')')
+        call me%add_str(xname//' = '//xcstr//' + '//rstr//' * np.outer(np.cos(u), np.sin(v))')
+        call me%add_str(yname//' = '//ycstr//' + '//rstr//' * np.outer(np.sin(u), np.sin(v))')
+        call me%add_str(zname//' = '//zcstr//' + '//rstr//' * np.outer(np.ones(np.size(u)), np.cos(v))')
+        call me%add_str('ax.plot_surface('//xname//', '//yname//', '//zname//extras//')')
+        call me%add_str('')
+
+    else
+        istat = -1
+        write(error_unit,'(A)') 'Error in add_sphere: pyplot class not properly initialized.'
+    end if
+
+    end subroutine add_sphere
+!*****************************************************************************************
+
 !*****************************************************************************************
 !> author: Jacob Williams
 !
@@ -762,6 +840,35 @@ subroutine integer_to_string(i, s)
     end subroutine integer_to_string
 !*****************************************************************************************
 
+!*****************************************************************************************
+!> author: Jacob Williams
+!
+! Real scalar to string.
+
+    subroutine real_to_string(v, fmt, str)
+
+    real(wp),                      intent(in)  :: v         !! real values
+    character(len=*),              intent(in)  :: fmt       !! real format string
+    character(len=:), allocatable, intent(out) :: str       !! real values stringified
+
+    integer                     :: istat     !! IO status
+    character(len=max_real_len) :: tmp       !! dummy string
+
+    if (fmt=='*') then
+        write(tmp, *, iostat=istat) v
+    else
+        write(tmp, fmt, iostat=istat) v
+    end if
+    if (istat/=0) then
+        write(error_unit,'(A)') 'Error in real_to_string'
+        str = '****'
+    else
+        str = trim(adjustl(tmp))
+    end if
+
+    end subroutine real_to_string
+!*****************************************************************************************
+
 !*****************************************************************************************
 !> author: Jacob Williams
 !
@@ -937,7 +1044,29 @@ subroutine finish_ops(me)
         call me%add_str('')
     end if
     if (me%axis_equal) then
-        call me%add_str('ax.axis("equal")')
+        if (me%mplot3d) then
+            call me%add_str('ax.set_aspect("equal")')
+            call me%add_str('')
+
+            call me%add_str('def set_axes_equal(ax):')
+            call me%add_str('    x_limits = ax.get_xlim3d()')
+            call me%add_str('    y_limits = ax.get_ylim3d()')
+            call me%add_str('    z_limits = ax.get_zlim3d()')
+            call me%add_str('    x_range = abs(x_limits[1] - x_limits[0])')
+            call me%add_str('    x_middle = np.mean(x_limits)')
+            call me%add_str('    y_range = abs(y_limits[1] - y_limits[0])')
+            call me%add_str('    y_middle = np.mean(y_limits)')
+            call me%add_str('    z_range = abs(z_limits[1] - z_limits[0])')
+            call me%add_str('    z_middle = np.mean(z_limits)')
+            call me%add_str('    plot_radius = 0.5*max([x_range, y_range, z_range])')
+            call me%add_str('    ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])')
+            call me%add_str('    ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])')
+            call me%add_str('    ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])')
+            call me%add_str('set_axes_equal(ax)')
+
+        else
+            call me%add_str('ax.axis("equal")')
+        end if
         call me%add_str('')
     end if
 

src/tests/test.f90

@@ -19,6 +19,7 @@ program test
     real(wp), dimension(n)   :: yerr  !! error values for bar chart
     real(wp), dimension(n)   :: sx    !! sin(x) values
     real(wp), dimension(n)   :: cx    !! cos(x) values
+    real(wp), dimension(n)   :: zx    !!
     real(wp), dimension(n)   :: tx    !! sin(x)*cos(x) values
     real(wp), dimension(n,n) :: z     !! z matrix for contour plot
     type(pyplot)             :: plt   !! pytplot handler
@@ -28,11 +29,15 @@ program test
     real(wp), dimension(n,n) :: mat   !! image values
     integer                  :: istat !! status code
 
+    real(wp),parameter :: pi = acos(-1.0_wp)
+    real(wp),parameter :: deg2rad = pi/180.0_wp
+
     !generate some data:
     x    = [(real(i,wp), i=0,size(x)-1)]/5.0_wp
     sx   = sin(x)
     cx   = cos(x)
     tx   = sx * cx
+    zx   = 0.0_wp
     yerr = abs(sx*.25_wp)
 
     do i=1,n
@@ -124,5 +129,28 @@ program test
                         dpi='200', &
                         transparent=.true.,istat=istat)
 
+    !3D plot:
+    call plt%initialize(grid=.true.,&
+                        xlabel='x (km)',ylabel='y (km)',zlabel='z (km)',&
+                        title='Orbit',&
+                        axis_equal=.true.,&
+                        mplot3d=.true.,&
+                        figsize=[20,10])
+
+    x    = [(real(i,wp), i=0,size(x)-1)]/5.0_wp
+    sx   = 7000.0_wp * cos(x * deg2rad)
+    cx   = 7000.0_wp * sin(x * deg2rad)
+    zx = 0.0_wp
+    call plt%add_3d_plot(sx,cx,zx,&
+                            label='orbit',linestyle='r-',&
+                            linewidth=2,istat=istat)
+    call plt%add_sphere(6378.0_wp,0.0_wp,0.0_wp,0.0_wp,&
+                        color='Blue',n_facets=500,&
+                        antialiased=.true.,istat=istat)
+    call plt%savefig('orbit.png', &
+                      pyfile='orbit.py', &
+                      dpi='200', &
+                      transparent=.true.,istat=istat)
+
     end program test
 !*****************************************************************************************