getdata -> getdata2d

ext/BatsrusMakieExt/BatsrusMakieExt.jl

@@ -1,7 +1,7 @@
 module BatsrusMakieExt
 
 using Batsrus, Printf
-import Batsrus: findindex, hasunit, getunit, getdata
+import Batsrus: findindex, hasunit, getunit, getdata2d
 import Batsrus.UnitfulBatsrus
 import Makie
 

ext/BatsrusMakieExt/typerecipe.jl

@@ -19,7 +19,7 @@ end
 "Conversion for 2D plots."
 function Makie.convert_arguments(P::Makie.GridBased, bd::BATLData, var::String;
    plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1)
-   x, y, w = getdata(bd, var, plotrange, plotinterval)
+   x, y, w = getdata2d(bd, var, plotrange, plotinterval)
 
    unitx = getunit(bd, bd.head.variables[1])
    unity = getunit(bd, bd.head.variables[2])

src/Batsrus.jl

@@ -5,7 +5,11 @@ module Batsrus
 
 using Printf, Requires
 
-export BATLData
+export BATLData,
+   load, readlogdata, readtecdata, showhead, # io
+   getvars, getvar, cutdata, subvolume, subsurface, # select
+   Batl, convertTECtoVTU, convertIDLtoVTK, readhead, readtree, getConnectivity, # vtk
+   getdata2d # plot/utility
 
 "Type for the file information."
 struct FileList

src/io.jl

@@ -1,7 +1,5 @@
 # All the IO related APIs.
 
-export load, readlogdata, readtecdata, showhead
-
 # Fortran binary format includes record end/start tags that needs to be skipped.
 # If there are continuous blocks, we usually skip 2*TAG between actual reading.
 const TAG = 4 # Fortran record tag size

src/plot/plots.jl

@@ -27,7 +27,7 @@ using RecipesBase
          x, y
       end
    elseif ndim == 2
-      x, y, w = getdata(bd, var, plotrange, plotinterval)
+      x, y, w = getdata2d(bd, var, plotrange, plotinterval)
 
       unitx = getunit(bd, bd.head.variables[1])
       unity = getunit(bd, bd.head.variables[2])

src/plot/pyplot.jl

@@ -470,7 +470,7 @@ Wrapper over `contour` in matplotlib.
 function PyPlot.contour(bd::BATLData, var::AbstractString, levels::Int=0; ax=nothing,
    plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, kwargs...)
 
-   Xi, Yi, Wi = getdata(bd, var, plotrange, plotinterval; innermask)
+   Xi, Yi, Wi = getdata2d(bd, var, plotrange, plotinterval; innermask)
 
    if isnothing(ax) ax = plt.gca() end
 
@@ -490,7 +490,7 @@ Wrapper over `contourf` in matplotlib.
 function PyPlot.contourf(bd::BATLData, var::AbstractString, levels::Int=0; ax=nothing,
    plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, kwargs...)
 
-   Xi, Yi, Wi = getdata(bd, var, plotrange, plotinterval; innermask)
+   Xi, Yi, Wi = getdata2d(bd, var, plotrange, plotinterval; innermask)
 
    if isnothing(ax) ax = plt.gca() end
 
@@ -569,7 +569,8 @@ Wrapper over `plot_surface` in matplotlib.
 function PyPlot.plot_surface(bd::BATLData, var::AbstractString;
    plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, kwargs...)
 
-   Xi, Yi, Wi = getdata(bd, var, plotrange, plotinterval; griddim=2, innermask)
+   xi, yi, Wi = getdata2d(bd, var, plotrange, plotinterval; innermask)
+   Xi, Yi = meshgrid(xi, yi)
 
    plot_surface(Xi, Yi, Wi; kwargs...)
 end
@@ -584,11 +585,11 @@ 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)
+   xi, yi, Wi = getdata2d(bd, var, plotrange, plotinterval; innermask)
 
    if isnothing(ax) ax = plt.gca() end
 
-   c = ax.pcolormesh(Xi, Yi, Wi; kwargs...)
+   c = ax.pcolormesh(xi, yi, Wi; kwargs...)
 end
 
 """

src/plot/utility.jl

@@ -1,8 +1,15 @@
 # Utility functions for plotting.
 
-"Prepare 2D data arrays for passing to plotting functions."
-function getdata(bd::BATLData, var::AbstractString, plotrange, plotinterval; griddim=1,
-   innermask::Bool=false)
+"""
+    getdata2d(bd::BATLData, var::AbstractString, plotrange=[-Inf, Inf, -Inf, Inf],
+       plotinterval=Inf; innermask=false)
+
+Return 2D slices of data `var` from `bd`. If `plotrange` is not set, output data resolution
+is the same as the original. If `innermask==true`, then the inner boundary cells are set to
+NaN.
+"""
+function getdata2d(bd::BATLData, var::AbstractString,
+   plotrange::Vector=[-Inf, Inf, -Inf, Inf], plotinterval::Real=Inf; innermask::Bool=false)
    x, w, ndim = bd.x, bd.w, bd.head.ndim
    @assert ndim == 2 "data must be in 2D!"
 
@@ -33,21 +40,23 @@ function getdata(bd::BATLData, var::AbstractString, plotrange, plotinterval; gri
       yrange = range(x[1,1,2], x[1,end,2], length=size(x,2))
       if all(isinf.(plotrange))
          xi, yi = xrange, yrange
-         Xi, Yi = meshgrid(xi, yi)
          Wi = w[:,:,varIndex_]' # Matplotlib does not accept view!
       else
-         if plotrange[1] == -Inf plotrange[1] = minimum(xrange) end
-         if plotrange[2] ==  Inf plotrange[2] = maximum(xrange) end
-         if plotrange[3] == -Inf plotrange[3] = minimum(yrange) end
-         if plotrange[4] ==  Inf plotrange[4] = maximum(yrange) end
+         if plotrange[1] == -Inf plotrange[1] = xrange[1] end
+         if plotrange[2] ==  Inf plotrange[2] = xrange[end] end
+         if plotrange[3] == -Inf plotrange[3] = yrange[1] end
+         if plotrange[4] ==  Inf plotrange[4] = yrange[end] end
 
-         xi = range(plotrange[1], stop=plotrange[2], step=plotinterval)
-         yi = range(plotrange[3], stop=plotrange[4], step=plotinterval)
+         if isinf(plotinterval)
+            xi = range(plotrange[1], stop=plotrange[2], step=xrange[2] - xrange[1])
+            yi = range(plotrange[3], stop=plotrange[4], step=yrange[2] - yrange[1])
+         else
+            xi = range(plotrange[1], stop=plotrange[2], step=plotinterval)
+            yi = range(plotrange[3], stop=plotrange[4], step=plotinterval)
+         end
 
-         spline = @views Spline2D(xrange, yrange, w[:,:,varIndex_])
-         Xi, Yi = meshgrid(xi, yi)
-         wi = @views spline(Xi[:], Yi[:])
-         Wi = reshape(wi, size(Xi))
+         interp = @views cubic_spline_interpolation((xrange, yrange), w[:,:,varIndex_])
+         Wi = [interp(i, j) for j in yi, i in xi]
       end
    end
 
@@ -56,18 +65,14 @@ function getdata(bd::BATLData, var::AbstractString, plotrange, plotinterval; gri
       varIndex_ = findlast(x->x=="rbody", bd.head.variables)
       isnothing(varIndex_) && error("rbody not found in file header parameters!")
       ParamIndex_ = varIndex_ - ndim - bd.head.nw
-      @inbounds for i = eachindex(Xi, Yi)
-         if Xi[i]^2 + Yi[i]^2 < bd.head.eqpar[ParamIndex_]^2
+      @inbounds for i in CartesianIndices(Wi)
+         if xi[i[1]]^2 + yi[i[2]]^2 < bd.head.eqpar[ParamIndex_]^2
             Wi[i] = NaN
          end
       end
    end
 
-   if griddim == 1
-      return xi, yi, Wi
-   else
-      return Xi, Yi, Wi
-   end
+   xi, yi, Wi
 end
 
 "Find variable index in the BATSRUS data."

src/select.jl

@@ -1,5 +1,4 @@
-export getvars, getvar, cutdata, subvolume, subsurface
-
+# Data manipulation.
 
 """
 	cutdata(data, var; plotrange=[-Inf,Inf,-Inf,Inf], dir="x", sequence=1)

src/vtk.jl

@@ -2,9 +2,6 @@
 
 using FortranFiles, WriteVTK, Glob, LightXML
 
-export convertTECtoVTU, convertIDLtoVTK, readhead, readtree, getConnectivity
-export Batl
-
 # Named indexes of iTree_IA
 const status_   =  Int8(1) # used, unused, to be refined, to be coarsened, etc.
 const level_    =  Int8(2) # grid AMR level for this node

test/runtests.jl

@@ -46,6 +46,9 @@ end
       @test bd.head.time == 25.6f0
       @test extrema(bd.x) == (-127.5f0, 127.5f0)
       @test extrema(bd.w) == (-0.79985905f0, 1.9399388f0)
+      plotrange = [-10.0, 10.0, -Inf, Inf]
+      x, y, w = Batsrus.getdata2d(bd, "rho", plotrange)
+      @test w[1,end] == 0.6848978549242021
    end
 
    @testset "Reading 2D unstructured binary" begin