Add pcolormesh

src/plot/pyplot.jl

@@ -4,7 +4,7 @@ using PyPlot
 using Dierckx: Spline2D
 
 export plotdata, plotlogdata, plot, scatter, contour, contourf, plot_surface,
-   tricontourf, plot_trisurf, streamplot, streamslice, quiver, cutplot
+   tricontourf, plot_trisurf, streamplot, streamslice, quiver, cutplot, pcolormesh
 
 @static if matplotlib.__version__ >= "3.3"
    matplotlib.rc("image", cmap="turbo") # set default colormap
@@ -644,6 +644,23 @@ function PyPlot.plot_surface(bd::BATLData, var::AbstractString;
    plot_surface(Xi, Yi, Wi'; kwargs...)
 end
 
+
+"""
+    pcolormesh(data, var, levels=0; ax=nothing, plotrange=[-Inf,Inf,-Inf,Inf],
+       plotinterval=0.1, innermask=false, kwargs...)
+
+Wrapper over `pcolormesh` in matplotlib.
+"""
+function PyPlot.pcolormesh(bd::BATLData, var::AbstractString; ax=nothing,
+   plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, kwargs...)
+
+   Xi, Yi, Wi = getdata(bd, var, plotrange, plotinterval; innermask)
+
+   if isnothing(ax) ax = plt.gca() end
+
+   c = ax.pcolormesh(Xi, Yi, Wi'; kwargs...)
+end
+
 """
     streamplot(data, var, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1,
        kwargs...)

test/runtests.jl

@@ -117,16 +117,16 @@ end
          @test plt isa Heatmap
       end
 
-      @testset "1D ascii" begin
+      @testset "PyPlot" begin
+         # 1D ascii
          file = "1d__raw_2_t25.60000_n00000258.out"
          bd = load(file, dir=datapath, verbose=false)
          plotdata(bd, "p", plotmode="line")
          line = get(gca().lines, 0)
          @test line.get_xdata() ≈ bd.x
          @test line.get_ydata() ≈ bd.w[:,10]
-      end
-
-      @testset "2D structured binary" begin
+
+         # 2D structured binary
          file = "z=0_raw_1_t25.60000_n00000258.out"
          bd = load(file, dir=datapath)
          plotdata(bd, "p bx;by", plotmode="contbar streamover")
@@ -140,15 +140,16 @@ end
          @test isa(gca(), PyPlot.PyObject)
          PyPlot.streamplot(bd, "bx;by")
          @test isa(gca(), PyPlot.PyObject)
+         p = PyPlot.pcolormesh(bd, "p").get_array()
+         @test p[end] == 0.1f0
          plt.close()
          fig = plt.figure()
          ax = fig.add_subplot(111, projection="3d")
          plot_surface(bd, "rho")
          @test isa(gca(), PyPlot.PyObject)
          plt.close()
-      end
 
-      @testset "2D AMR Cartesian" begin
+         # 2D AMR Cartesian
          file = "bx0_mhd_6_t00000100_n00000352.out"
          bd = load(file, dir=datapath)
          plotdata(bd, "P", plotmode="contbar")