Specify return type (#64)

* Specify return type

* Drop getvars

* Add tests and benchmarks

* Bump patch

Project.toml

@@ -1,7 +1,7 @@
 name = "Batsrus"
 uuid = "e74ebddf-6ac1-4047-a0e5-c32c99e57753"
 authors = ["Hongyang Zhou <[email protected]>"]
-version = "0.6.8"
+version = "0.6.9"
 
 [deps]
 FortranFiles = "c58ffaec-ab22-586d-bfc5-781a99fd0b10"

benchmark/benchmarks.jl

@@ -19,6 +19,9 @@ file = joinpath(directory, files[1])
 SUITE["read"]["ASCII"] = @benchmarkable load($file)
 
 file = joinpath(directory, files[2])
+bd = load(file)
+SUITE["read"]["Extract density"] = @benchmarkable Batsrus.getvar($bd, "Rho")
+SUITE["read"]["Extract B"] = @benchmarkable Batsrus.getvar($bd, "B")
 SUITE["read"]["Binary structured"] = @benchmarkable load($file)
 
 file = joinpath(directory, files[3])

docs/src/man/manual.md

@@ -40,13 +40,6 @@ get_var_range(bd, "rho")
 bd["rho"]
 ```
 
-- Derived variables
-
-```julia
-v = getvars(bd, ["Bx", "By", "Bz"])
-Bmag = bd["Bmag"]
-```
-
 - Extracting data at a given location
 
 ```julia
@@ -62,6 +55,12 @@ point2 = Float32[10.0, 1.0]
 w = interp1d(bd, "rho", point1, point2)
 ```
 
+- Derived variables
+
+```julia
+Bmag = bd["Bmag"]
+```
+
 Here is a full list of predefined derived quantities:
 
 | Derived variable name | Meaning                          | Required variable |

src/Batsrus.jl

@@ -12,7 +12,7 @@ using StaticArrays: SVector, @SMatrix, SA
 
 export BATLData,
    load, readlogdata, readtecdata, showhead, # io
-   getvars, getvar, cutdata, subvolume, subsurface, # select
+   getvar, cutdata, subvolume, subsurface, # select
    Batl, convertTECtoVTU, convertIDLtoVTK, readhead, readtree, getConnectivity, # vtk
    interp1d, interp2d, slice1d, get_var_range, squeeze, get_range # plot/utility
 
@@ -32,14 +32,28 @@ struct FileList
    lenhead::Int
 end
 
+struct BATLHead
+   ndim::Int32
+   headline::SubString{String}
+   it::Int32
+   time::Float32
+   gencoord::Bool
+   neqpar::Int32
+   nw::Int32
+   nx::Vector{Int32}
+   eqpar::Vector{Float32}
+   variables::Vector{SubString{String}}
+   wnames::Vector{SubString{String}}
+end
+
 "Primary Batsrus data storage type."
-struct BATLData{dim, T<:AbstractFloat}
+struct BATLData{dim, T<:AbstractFloat, U}
    "header information"
-   head::NamedTuple
+   head::BATLHead
    "grid"
-   x::Array{T}
+   x::U
    "variables"
-   w::Array{T}
+   w::U
    "file information"
    list::FileList
 end

src/io.jl

@@ -20,14 +20,14 @@ function load(file::AbstractString; npict::Int=1, verbose::Bool=false)
    end
    seekstart(fileID) # Rewind to start
 
-   ## Read data from files
    # Skip npict-1 snapshots (since we only want the npict-th snapshot)
    skip(fileID, pictsize*(npict-1))
 
    filehead = getfilehead(fileID, filelist)
    # Read data
    if filelist.type == :ascii
-      x, w = allocateBuffer(filehead, Float64) # why Float64?
+      T = Float64 # why Float64?
+      x, w = allocateBuffer(filehead, T)
       getascii!(x, w, fileID)
    else
       skip(fileID, TAG) # skip record start tag
@@ -40,7 +40,7 @@ function load(file::AbstractString; npict::Int=1, verbose::Bool=false)
 
    #setunits(filehead,"")
 
-	data = BATLData{Int(filehead.ndim), eltype(w)}(filehead, x, w, filelist)
+	data = BATLData{Int(filehead.ndim), T, typeof(w)}(filehead, x, w, filelist)
 end
 
 "Read information from log file."
@@ -313,9 +313,8 @@ function getfilehead(fileID::IOStream, filelist::FileList)
 	# Produce a wnames from the last file
    wnames = variables[ndim+1:ndim+nw]
 
-   head = (ndim=ndim, headline=headline, it=it, time=t, gencoord=gencoord,
-		neqpar=neqpar, nw=nw, nx=nx, eqpar=eqpar, variables=variables,
-		wnames=wnames)
+   head = BATLHead(ndim, headline, it, t, gencoord,
+		neqpar, nw, nx, eqpar, variables, wnames)
 end
 
 function skipline(s::IO)
@@ -381,7 +380,7 @@ function getfilesize(fileID::IOStream, type::Symbol, lenstr::Int32)
 end
 
 "Create buffer for x and w."
-function allocateBuffer(filehead::NamedTuple, T::DataType)
+function allocateBuffer(filehead::BATLHead, T::DataType)
    if filehead.ndim == 1
       n1 = filehead.nx[1]
       x  = Array{T,2}(undef, n1, filehead.ndim)

src/select.jl

@@ -163,16 +163,14 @@ function subvolume(x, y, z, u, v, w, limits)
 end
 
 """
-    getvars(bd::BATLData, var::AbstractString) -> Array
+    getvar(bd::BATLData, var::AbstractString) -> Array
 
 Return variable data from string `var`. This is also supported via direct indexing,
 
 # Examples
 ```
 bd["rho"]
 ```
-
-See also: [`getvars`](@ref).
 """
 function getvar(bd::BATLData{ndim, T}, var::AbstractString) where {ndim, T}
    if var in keys(variables_predefined)
@@ -189,26 +187,11 @@ end
 @inline @Base.propagate_inbounds Base.getindex(bd::BATLData, var::AbstractString) =
    getvar(bd, var)
 
-"""
-    getvars(bd::BATLData, names::Vector) -> Dict
-
-Return variables' data as a dictionary from vector of `names`.
-See also: [`getvar`](@ref).
-"""
-function getvars(bd::BATLData{ndim, U}, names::Vector{T}) where {ndim, U, T<:AbstractString}
-   dict = Dict{T, Array{U}}()
-   for name in names
-      dict[name] = getvar(bd, name)
-   end
 
-   dict
-end
 
 "Construct vectors from scalar components."
-function _fill_vector_from_scalars(bd::BATLData{ndim, T}, vstr1, vstr2, vstr3) where {ndim, T}
-   v1 = getvar(bd, vstr1)
-   v2 = getvar(bd, vstr2)
-   v3 = getvar(bd, vstr3)
+function _fill_vector_from_scalars(bd::BATLData{ndim, T, U}, var) where {ndim, T, U}
+   v1, v2, v3 = _get_vectors(bd, var)
    v = Array{T, ndims(v1)+1}(undef, 3, size(v1)...)
 
    Rpost = CartesianIndices(size(v1))
@@ -221,7 +204,55 @@ function _fill_vector_from_scalars(bd::BATLData{ndim, T}, vstr1, vstr2, vstr3) w
    v
 end
 
-function _get_anisotropy(bd::BATLData{2, T}, species=0) where {T}
+function _get_magnitude2(bd::BATLData{2, T, U}, var=:B) where {T, U}
+   vx, vy, vz = _get_vectors(bd, var)
+   v = similar(vx)::Array{T, 2}
+
+   @simd for i in eachindex(v)
+      v[i] = vx[i]^2 + vy[i]^2 + vz[i]^2
+   end
+
+   v
+end
+
+function _get_magnitude(bd::BATLData{2, T, U}, var=:B) where {T, U}
+   vx, vy, vz = _get_vectors(bd, var)
+   v = similar(vx)::Array{T, 2}
+
+   @simd for i in eachindex(v)
+      v[i] = √(vx[i]^2 + vy[i]^2 + vz[i]^2)
+   end
+
+   v
+end
+
+function _get_vectors(bd::BATLData, var)
+   if var == :B
+      vx = bd["Bx"]
+      vy = bd["By"]
+      vz = bd["Bz"]
+   elseif var == :E
+      vx = bd["Ex"]
+      vy = bd["Ey"]
+      vz = bd["Ez"]
+   elseif var == :U
+      vx = bd["Ux"]
+      vy = bd["Uy"]
+      vz = bd["Uz"]
+   elseif var == :U0
+      vx = bd["UxS0"]
+      vy = bd["UyS0"]
+      vz = bd["UzS0"]
+   elseif var == :U1
+      vx = bd["UxS1"]
+      vy = bd["UyS1"]
+      vz = bd["UzS1"] 
+   end
+
+   vx, vy, vz
+end
+
+function _get_anisotropy(bd::BATLData{2, T, U}, species=0) where {T, U}
    Bx, By, Bz = bd["Bx"], bd["By"], bd["Bz"]
    # Rotate the pressure tensor to align the 3rd direction with B
    pop = string(species)
@@ -232,7 +263,7 @@ function _get_anisotropy(bd::BATLData{2, T}, species=0) where {T}
    Pxz = bd["pXZS"*pop]
    Pyz = bd["pYZS"*pop]
 
-   Paniso = similar(Pxx)
+   Paniso = similar(Pxx)::Array{T, 2}
 
    @inbounds for j in axes(Pxx, 2), i in axes(Pxx, 1)  
       b̂ = normalize(SA[Bx[i,j], By[i,j], Bz[i,j]])
@@ -249,18 +280,18 @@ function _get_anisotropy(bd::BATLData{2, T}, species=0) where {T}
 end
 
 # Define derived parameters
-const variables_predefined = Dict(
-   "B2" => (bd -> @. $getvar(bd, "Bx")^2 + $getvar(bd, "By")^2 + $getvar(bd, "Bz")^2),
-   "E2" => (bd -> @. $getvar(bd, "Ex")^2 + $getvar(bd, "Ey")^2 + $getvar(bd, "Ez")^2),
-   "U2" => (bd -> @. $getvar(bd, "Ux")^2 + $getvar(bd, "Uy")^2 + $getvar(bd, "Uz")^2),
-   "Ue2" => (bd -> @. $getvar(bd, "uxS0")^2 + $getvar(bd, "uyS0")^2 + $getvar(bd, "uzS0")^2),
-   "Ui2" => (bd -> @. $getvar(bd, "uxS1")^2 + $getvar(bd, "uyS1")^2 + $getvar(bd, "uzS1")^2),
-   "Bmag" => (bd -> @. sqrt($getvar(bd, "B2"))),
-   "Emag" => (bd -> @. sqrt($getvar(bd, "E2"))),
-   "Umag" => (bd -> @. sqrt($getvar(bd, "U2"))),
-   "B" => (bd -> _fill_vector_from_scalars(bd, "Bx", "By", "Bz")),
-   "E" => (bd -> _fill_vector_from_scalars(bd, "Ex", "Ey", "Ez")),
-   "U" => (bd -> _fill_vector_from_scalars(bd, "Ux", "Uy", "Uz")),
+const variables_predefined = Dict{String, Function}(
+   "B2" => (bd -> _get_magnitude2(bd, :B)),
+   "E2" => (bd -> _get_magnitude2(bd, :E)),
+   "U2" => (bd -> _get_magnitude2(bd, :U)),
+   "Ue2" => (bd -> _get_magnitude2(bd, :U0)),
+   "Ui2" => (bd -> _get_magnitude2(bd, :U1)),
+   "Bmag" => (bd -> _get_magnitude(bd, :B)),
+   "Emag" => (bd -> _get_magnitude(bd, :E)),
+   "Umag" => (bd -> _get_magnitude(bd, :U)),
+   "B" => (bd -> _fill_vector_from_scalars(bd, :B)),
+   "E" => (bd -> _fill_vector_from_scalars(bd, :E)),
+   "U" => (bd -> _fill_vector_from_scalars(bd, :U)),
    "Anisotropy0" => (bd -> _get_anisotropy(bd, 0)),
    "Anisotropy1" => (bd -> _get_anisotropy(bd, 1)),
 )

src/utility.jl

@@ -12,9 +12,9 @@ data resolution is the same as the original.
 - `rbody=1.0`: Radius of the inner mask. Used when the rbody parameter is not found in the header.
 - `useMatplotlib=true`: Whether to Matplotlib (somehow faster) or NaturalNeighbours for scattered interpolation.
 """
-function interp2d(bd::BATLData{2, T}, var::AbstractString,
+function interp2d(bd::BATLData{2, T, U}, var::AbstractString,
    plotrange::Vector=[-Inf, Inf, -Inf, Inf], plotinterval::Real=Inf;
-   innermask::Bool=false, rbody::Real=1.0, useMatplotlib::Bool=true) where {T}
+   innermask::Bool=false, rbody::Real=1.0, useMatplotlib::Bool=true) where {T, U}
    x, w = bd.x, bd.w
    varIndex_ = findindex(bd, var)
 
@@ -133,7 +133,7 @@ end
 
 Interpolate `var` at spatial point `loc` in `bd`.
 """
-function interp1d(bd::BATLData{2, T}, var::AbstractString, loc::AbstractVector{<:AbstractFloat}) where {T}
+function interp1d(bd::BATLData{2, T, U}, var::AbstractString, loc::AbstractVector{<:AbstractFloat}) where {T, U}
    @assert !bd.head.gencoord "Only accept structured grids!"
 
    x = bd.x
@@ -150,7 +150,7 @@ end
 
 Interpolate `var` along a line from `point1` to `point2` in `bd`.
 """
-function interp1d(bd::BATLData{2, T}, var::AbstractString, point1::Vector, point2::Vector) where {T}
+function interp1d(bd::BATLData{2, T, U}, var::AbstractString, point1::Vector, point2::Vector) where {T, U}
    @assert !bd.head.gencoord "Only accept structured grids!"
 
    x = bd.x
@@ -179,13 +179,13 @@ function slice1d(bd, var, icut::Int=1, dir::Int=2)
 end
 
 "Return view of variable `var` in `bd`."
-function getview(bd::BATLData{1, T}, var) where T
+function getview(bd::BATLData{1, T, U}, var) where {T, U}
    varIndex_ = findindex(bd, var)
 
    v = @view bd.w[:,varIndex_]
 end
 
-function getview(bd::BATLData{2, T}, var) where T
+function getview(bd::BATLData{2, T, U}, var) where {T, U}
    varIndex_ = findindex(bd, var)
 
    v = @view bd.w[:,:,varIndex_]

test/runtests.jl

@@ -38,7 +38,6 @@ end
       end
       @test startswith(repr(bd), "filename : 1d")
       @test Batsrus.setunits(bd.head, "NORMALIZED")
-      @test isa(bd.head, NamedTuple)
       @test extrema(bd.x) == (-127.5, 127.5)
       @test extrema(bd.w) == (-0.79960780498, 1.9394335293)
    end
@@ -53,6 +52,7 @@ end
       x, y, w = interp2d(bd, "rho", plotrange)
       @test w[1,end] == 0.6848635077476501
       @test bd["B"][:,end,end] == Float32[1.118034, -0.559017, 0.0]
+      @test bd["Bmag"][128,2] == 0.9223745f0
       # Linear interpolation at a given point
       d = interp1d(bd, "rho", Float32[0.0, 0.0])
       @test d == 0.6936918f0
@@ -77,6 +77,7 @@ end
       plotrange = [-Inf, Inf, -Inf, Inf]
       x, y, w = interp2d(bd, "rho", plotrange, useMatplotlib=false)
       @test w[1,2] == 5.000018304080387
+      @test bd["Umag"][2] == 71.85452748407637
    end
 
    @testset "Reading 3D structured binary" begin
@@ -86,8 +87,6 @@ end
       X, Z, p = cutdata(bd, "p"; dir="y", sequence=1, plotrange)
       @test p[1] ≈ 0.560976f0 && p[2] ≈ 0.53704995f0
       @test size(bd["p"]) == (8,8,8)
-      vars = getvars(bd, ["p"])
-      @test size(vars["p"]) == (8,8,8)
    end
    #TODO: add tecplot tests
    #@testset "Reading Tecplot" begin