From 80e96e14c63a52253d70d8c749dec4aceb4d3036 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Sun, 10 Nov 2024 14:47:40 -1000 Subject: [PATCH] Primitive struct optimization (#66) * Use Enum for FileTypes; dispatch on FileType values * Work around derive type construction for BATS, reordering * Reduce allocations for lowercase check * Add type annotations * Int32 -> Int64 parameters * Add anisotropy benchmark * Bump minor version --- Project.toml | 2 +- benchmark/benchmarks.jl | 5 +- .../examples/visualization/demo_animate_2d.md | 5 +- docs/src/man/manual.md | 4 +- examples/vdf.jl | 6 +- ext/BatsrusMakieExt/typerecipe.jl | 4 +- src/Batsrus.jl | 45 +--- src/io.jl | 235 +++++++++--------- src/plot/plots.jl | 4 +- src/plot/pyplot.jl | 49 ++-- src/select.jl | 69 +++-- src/type.jl | 49 ++++ src/unit/UnitfulBatsrus.jl | 2 +- src/utility.jl | 28 +-- test/runtests.jl | 10 +- 15 files changed, 270 insertions(+), 247 deletions(-) create mode 100644 src/type.jl diff --git a/Project.toml b/Project.toml index b6be0ac4..1e383171 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Batsrus" uuid = "e74ebddf-6ac1-4047-a0e5-c32c99e57753" authors = ["Hongyang Zhou "] -version = "0.6.9" +version = "0.7.0" [deps] FortranFiles = "c58ffaec-ab22-586d-bfc5-781a99fd0b10" diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index cbf7dee2..ffc17e32 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -9,8 +9,7 @@ println() directory = artifact"testdata" files = ("1d__raw_2_t25.60000_n00000258.out", "z=0_raw_1_t25.60000_n00000258.out", - "bx0_mhd_6_t00000100_n00000352.out") - + "z=0_fluid_region0_0_t00001640_n00010142.out") const SUITE = BenchmarkGroup() @@ -25,4 +24,4 @@ SUITE["read"]["Extract B"] = @benchmarkable Batsrus.getvar($bd, "B") SUITE["read"]["Binary structured"] = @benchmarkable load($file) file = joinpath(directory, files[3]) -SUITE["read"]["Binary unstructured"] = @benchmarkable load($file) +SUITE["read"]["Anisotropy"] = @benchmarkable bd["Anisotropy1"] diff --git a/docs/examples/visualization/demo_animate_2d.md b/docs/examples/visualization/demo_animate_2d.md index 610d4273..03ab60cc 100644 --- a/docs/examples/visualization/demo_animate_2d.md +++ b/docs/examples/visualization/demo_animate_2d.md @@ -190,9 +190,8 @@ function get_vector(bd, var, plotrange, plotinterval) X, Y = eachslice(x, dims=3) X, Y = vec(X), vec(Y) varstream = split(var, ";") - wnames = lowercase.(bd.head.wnames) - var1_ = findfirst(x->x==lowercase(varstream[1]), wnames) - var2_ = findfirst(x->x==lowercase(varstream[2]), wnames) + var1_ = findfirst(x->lowercase(x)==lowercase(varstream[1]), bd.head.wnames) + var2_ = findfirst(x->lowercase(x)==lowercase(varstream[2]), bd.head.wnames) # Create grid values first. xi = range(Float64(plotrange[1]), stop=Float64(plotrange[2]), step=plotinterval) yi = range(Float64(plotrange[3]), stop=Float64(plotrange[4]), step=plotinterval) diff --git a/docs/src/man/manual.md b/docs/src/man/manual.md index 08d25d8f..b9d15d23 100644 --- a/docs/src/man/manual.md +++ b/docs/src/man/manual.md @@ -71,10 +71,12 @@ Here is a full list of predefined derived quantities: | Bmag | Magnetic field magnitude | Bx, By, Bz | | Emag | Electric field magnitude | Ex, Ey, Ez | | Umag | Velocity magnitude | Ux, Uy, Uz | +| Uemag | Electron Velocity magnitude | UxS0, UyS0, UzS0 | +| Uimag | Ion Velocity magnitude | UxS1, UyS1, UzS1 | | B | Magnetic field vector | Bx, By, Bz | | E | Electric field vector | Ex, Ey, Ez | | U | Velocity vector | Ux, Uy, Uz | -| Anisotropy | Pressure/Temperature anisotropy | B, P tensor | +| Anisotropy[0-1] | Pressure/Temperature anisotropy | B, P tensor | ## Output format conversion diff --git a/examples/vdf.jl b/examples/vdf.jl index 48b43f4b..aa416695 100644 --- a/examples/vdf.jl +++ b/examples/vdf.jl @@ -587,9 +587,9 @@ function HF_velocity() y = data.x[:,:,:,2] z = data.x[:,:,:,3] - ux_ = findfirst(x->x=="uxs1", lowercase.(data.head.wnames)) - uy_ = findfirst(x->x=="uys1", lowercase.(data.head.wnames)) - uz_ = findfirst(x->x=="uzs1", lowercase.(data.head.wnames)) + ux_ = findfirst(x->lowercase(x)=="uxs1", data.head.wnames) + uy_ = findfirst(x->lowercase(x)=="uys1", data.head.wnames) + uz_ = findfirst(x->lowercase(x)=="uzs1", data.head.wnames) Ux = @view data.w[:,:,:,ux_] Uy = @view data.w[:,:,:,uy_] diff --git a/ext/BatsrusMakieExt/typerecipe.jl b/ext/BatsrusMakieExt/typerecipe.jl index e27733ef..6699a499 100644 --- a/ext/BatsrusMakieExt/typerecipe.jl +++ b/ext/BatsrusMakieExt/typerecipe.jl @@ -1,7 +1,7 @@ # Type conversion from Batsrus to Makie "Conversion for 1D plots" -function Makie.convert_arguments(P::Makie.PointBased, bd::BATLData, var::String) +function Makie.convert_arguments(P::Makie.PointBased, bd::BATS, var::String) var_ = findindex(bd, var) if hasunit(bd) unitx = getunit(bd, bd.head.variables[1]) @@ -17,7 +17,7 @@ function Makie.convert_arguments(P::Makie.PointBased, bd::BATLData, var::String) end "Conversion for 2D plots." -function Makie.convert_arguments(P::Makie.GridBased, bd::BATLData, var::String; +function Makie.convert_arguments(P::Makie.GridBased, bd::BATS, var::String; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1) x, y, w = interp2d(bd, var, plotrange, plotinterval) diff --git a/src/Batsrus.jl b/src/Batsrus.jl index 6d04ddc8..bc704ebc 100644 --- a/src/Batsrus.jl +++ b/src/Batsrus.jl @@ -10,54 +10,13 @@ using Interpolations: cubic_spline_interpolation, BSpline, Linear, scale, interp import NaturalNeighbours as NN using StaticArrays: SVector, @SMatrix, SA -export BATLData, +export BATS, load, readlogdata, readtecdata, showhead, # io getvar, cutdata, subvolume, subsurface, # select Batl, convertTECtoVTU, convertIDLtoVTK, readhead, readtree, getConnectivity, # vtk interp1d, interp2d, slice1d, get_var_range, squeeze, get_range # plot/utility -"Type for the file information." -struct FileList - "filename" - name::String - "file type" - type::Symbol - "directory" - dir::String - "file size" - bytes::Int - "number of snapshots" - npictinfiles::Int - "length of meta data" - 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, U} - "header information" - head::BATLHead - "grid" - x::U - "variables" - w::U - "file information" - list::FileList -end - +include("type.jl") include("unit/UnitfulBatsrus.jl") using .UnitfulBatsrus diff --git a/src/io.jl b/src/io.jl index 939ce3b2..7db19d86 100644 --- a/src/io.jl +++ b/src/io.jl @@ -7,8 +7,7 @@ const TAG = 4 # Fortran record tag size """ load(filename; npict=1, verbose=false) -Read BATSRUS output files. Stores the `npict` snapshot from an ascii or binary data file -into the arrays of coordinates `x` and data `w`. +Read BATSRUS output files. Stores the `npict` snapshot from an ascii or binary data file into the arrays of coordinates `x` and data `w`. """ function load(file::AbstractString; npict::Int=1, verbose::Bool=false) filelist, fileID, pictsize = getfiletype(file) @@ -19,28 +18,27 @@ function load(file::AbstractString; npict::Int=1, verbose::Bool=false) throw(ArgumentError("Select snapshot $npict out of range $(filelist.npictinfiles)!")) end seekstart(fileID) # Rewind to start - - # Skip npict-1 snapshots (since we only want the npict-th snapshot) + # Jump to the npict-th snapshot skip(fileID, pictsize*(npict-1)) - filehead = getfilehead(fileID, filelist) + head = getfilehead(fileID, filelist) # Read data - if filelist.type == :ascii + if filelist.type == AsciiBat T = Float64 # why Float64? - x, w = allocateBuffer(filehead, T) + x, w = allocateBuffer(head, T) getascii!(x, w, fileID) else skip(fileID, TAG) # skip record start tag - T = filelist.type == :real4 ? Float32 : Float64 - x, w = allocateBuffer(filehead, T) + T = filelist.type == Real4Bat ? Float32 : Float64 + x, w = allocateBuffer(head, T) getbinary!(x, w, fileID) end close(fileID) - #setunits(filehead,"") + #setunits(head,"") - data = BATLData{Int(filehead.ndim), T, typeof(w)}(filehead, x, w, filelist) + BATS(head, filelist, x, w) end "Read information from log file." @@ -74,8 +72,7 @@ end """ readtecdata(file; verbose=false) -Return header, data and connectivity from BATSRUS Tecplot outputs. Both 2D and -3D binary and ASCII formats are supported. +Return header, data and connectivity from BATSRUS Tecplot outputs. Both 2D and 3D binary and ASCII formats are supported. # Examples ``` file = "3d_ascii.dat" @@ -215,33 +212,33 @@ function getfiletype(file::AbstractString) # Check the appendix of file names if occursin(r"^.*\.(log)$", file) - type = :log + type = LogBat npictinfiles = 1 elseif occursin(r"^.*\.(dat)$", file) # Tecplot ascii format - type = :dat + type = TecBat npictinfiles = 1 else # Obtain filetype based on the length info in the first 4 bytes (Gabor's trick) lenhead = read(fileID, Int32) if lenhead != 79 && lenhead != 500 - type = :ascii + type = AsciiBat else # The length of the 2nd line decides between real4 & real8 # since it contains the time; which is real*8 | real*4 skip(fileID, lenhead + TAG) len = read(fileID, Int32) if len == 20 - type = :real4 + type = Real4Bat elseif len == 24 - type = :binary + type = Real8Bat else throw(ArgumentError("Incorrect formatted file: $file")) end end # Obtain file size & number of snapshots seekstart(fileID) - pictsize = getfilesize(fileID, type, lenhead) + pictsize = getfilesize(fileID, lenhead, Val(type))::Int npictinfiles = bytes ÷ pictsize end @@ -261,32 +258,29 @@ Obtain the header information from BATSRUS output file of `type` linked to `file function getfilehead(fileID::IOStream, filelist::FileList) type, lenstr = filelist.type, filelist.lenhead - ## Read header - pointer0 = position(fileID) - - if type == :ascii + if type == AsciiBat headline = readline(fileID) line = readline(fileID) |> split - it = Parsers.parse(Int, line[1]) - t = Parsers.parse(Float64, line[2]) - ndim = Parsers.parse(Int8, line[3]) + it = Parsers.parse(Int32, line[1]) + t = Parsers.parse(Float32, line[2]) + ndim = Parsers.parse(Int32, line[3]) neqpar = Parsers.parse(Int32, line[4]) - nw = Parsers.parse(Int8, line[5]) + nw = Parsers.parse(Int32, line[5]) gencoord = ndim < 0 ndim = abs(ndim) - nx = Parsers.parse.(Int64, split(readline(fileID))) + nx = Parsers.parse.(Int32, split(readline(fileID))) if neqpar > 0 - eqpar = Parsers.parse.(Float64, split(readline(fileID))) + eqpar = Parsers.parse.(Float32, split(readline(fileID))) end varname = readline(fileID) - elseif type ∈ (:real4, :binary) + elseif type ∈ (Real4Bat, Real8Bat) skip(fileID, TAG) headline = rstrip(String(read(fileID, lenstr))) skip(fileID, 2*TAG) it = read(fileID, Int32) t = read(fileID, Float32) ndim = read(fileID, Int32) - gencoord = (ndim < 0) + gencoord = ndim < 0 ndim = abs(ndim) neqpar = read(fileID, Int32) nw = read(fileID, Int32) @@ -303,18 +297,13 @@ function getfilehead(fileID::IOStream, filelist::FileList) skip(fileID, TAG) end - # Header length - pointer1 = position(fileID) - headlen = pointer1 - pointer0 - - # Set variables array - variables = split(varname) # returns a string array - - # Produce a wnames from the last file - wnames = variables[ndim+1:ndim+nw] + # Obtain output array + variables = split(varname) + # Obtain variable names + wnames = @view variables[ndim+1:ndim+nw] - head = BATLHead(ndim, headline, it, t, gencoord, - neqpar, nw, nx, eqpar, variables, wnames) + head = BatsHead(ndim, headline, it, t, gencoord, + neqpar, nw, nx, eqpar, variables, wnames) end function skipline(s::IO) @@ -326,73 +315,97 @@ function skipline(s::IO) return end -"Return the size in bytes for one snapshot." -function getfilesize(fileID::IOStream, type::Symbol, lenstr::Int32) - # Read header - pointer0 = position(fileID) +""" + getfilesize(fileID::IOStream, lenstr::Int32, ::Val{FileType}) - if type == :ascii - skipline(fileID) - line = readline(fileID) - line = split(line) - ndim = Parsers.parse(Int32, line[3]) - neqpar = Parsers.parse(Int32, line[4]) - nw = Parsers.parse(Int8, line[5]) - gencoord = ndim < 0 - ndim = abs(ndim) - nx = Parsers.parse.(Int64, split(readline(fileID))) - neqpar > 0 && skipline(fileID) - skipline(fileID) - elseif type ∈ (:real4, :binary) - skip(fileID, TAG + lenstr + 2*TAG + sizeof(Int32) + sizeof(Float32)) - ndim = abs(read(fileID, Int32)) - tmp = read(fileID, Int32) - nw = read(fileID, Int32) +Return the size in bytes for one snapshot. +""" +function getfilesize(fileID::IOStream, lenstr::Int32, ::Val{Real4Bat}) + pointer0 = position(fileID) # Record header start location + + skip(fileID, TAG + lenstr + 2*TAG + sizeof(Int32) + sizeof(Float32)) + ndim = abs(read(fileID, Int32)) + nt = read(fileID, Int32) + nw = read(fileID, Int32) + skip(fileID, 2*TAG) + nx = Vector{Int32}(undef, ndim) + read!(fileID, nx) + skip(fileID, 2*TAG) + if nt > 0 + tmp = zeros(Float32, nt) + read!(fileID, tmp) skip(fileID, 2*TAG) - nx = zeros(Int32, ndim) - read!(fileID, nx) + end + read(fileID, lenstr) + skip(fileID, TAG) + + pointer1 = position(fileID) + headlen = pointer1 - pointer0 # header length + # Calculate the snapshot size = header + data + recordmarks + pictsize = headlen + 8*(1 + nw) + 4*(ndim + nw)*prod(nx) +end + +function getfilesize(fileID::IOStream, lenstr::Int32, ::Val{Real8Bat}) + pointer0 = position(fileID) # Record header start location + + skip(fileID, TAG + lenstr + 2*TAG + sizeof(Int32) + sizeof(Float32)) + ndim = abs(read(fileID, Int32)) + nt = read(fileID, Int32) + nw = read(fileID, Int32) + skip(fileID, 2*TAG) + nx = Vector{Int32}(undef, ndim) + read!(fileID, nx) + skip(fileID, 2*TAG) + if nt > 0 + tmp = zeros(Float32, nt) + read!(fileID, tmp) skip(fileID, 2*TAG) - if tmp > 0 - tmp2 = zeros(Float32, tmp) - read!(fileID, tmp2) - skip(fileID, 2*TAG) - end - read(fileID, lenstr) - skip(fileID, TAG) end - # Header length + read(fileID, lenstr) + skip(fileID, TAG) + pointer1 = position(fileID) - headlen = pointer1 - pointer0 + headlen = pointer1 - pointer0 # header length # Calculate the snapshot size = header + data + recordmarks - nxs = prod(nx) - pictsize = - if type == :log - 1 - elseif type == :ascii - headlen + (18*(ndim + nw) + 1)*nxs - elseif type == :real4 - headlen + 8*(1 + nw) + 4*(ndim + nw)*nxs - elseif type == :binary - headlen + 8*(1 + nw) + 8*(ndim + nw)*nxs - end + pictsize = headlen + 8*(1 + nw) + 8*(ndim + nw)*prod(nx) +end + +function getfilesize(fileID::IOStream, lenstr::Int32, ::Val{AsciiBat}) + pointer0 = position(fileID) # Record header start location - pictsize + skipline(fileID) + line = readline(fileID) + line = split(line) + ndim = Parsers.parse(Int32, line[3]) + neqpar = Parsers.parse(Int32, line[4]) + nw = Parsers.parse(Int32, line[5]) + ndim = abs(ndim) + nx = Parsers.parse.(Int64, split(readline(fileID))) + neqpar > 0 && skipline(fileID) + skipline(fileID) + + pointer1 = position(fileID) + headlen = pointer1 - pointer0 # header length + # Calculate the snapshot size = header + data + recordmarks + pictsize = headlen + (18*(ndim + nw) + 1)*prod(nx) end +getfilesize(fileID::IOStream, lenstr::Int32, ::Val{LogBat}) = 1 + "Create buffer for x and w." -function allocateBuffer(filehead::BATLHead, T::DataType) - if filehead.ndim == 1 - n1 = filehead.nx[1] - x = Array{T,2}(undef, n1, filehead.ndim) - w = Array{T,2}(undef, n1, filehead.nw) - elseif filehead.ndim == 2 - n1, n2 = filehead.nx - x = Array{T,3}(undef, n1, n2, filehead.ndim) - w = Array{T,3}(undef, n1, n2, filehead.nw) - elseif filehead.ndim == 3 - n1, n2, n3 = filehead.nx - x = Array{T,4}(undef, n1, n2, n3, filehead.ndim) - w = Array{T,4}(undef, n1, n2, n3, filehead.nw) +function allocateBuffer(head::BatsHead, T::DataType) + if head.ndim == 1 + n1 = head.nx[1] + x = Array{T,2}(undef, n1, head.ndim) + w = Array{T,2}(undef, n1, head.nw) + elseif head.ndim == 2 + n1, n2 = head.nx + x = Array{T,3}(undef, n1, n2, head.ndim) + w = Array{T,3}(undef, n1, n2, head.nw) + elseif head.ndim == 3 + n1, n2, n3 = head.nx + x = Array{T,4}(undef, n1, n2, n3, head.ndim) + w = Array{T,4}(undef, n1, n2, n3, head.nw) end x, w @@ -455,7 +468,7 @@ function getbinary!(x::Array{T, 4}, w, fileID::IOStream) where T end """ - setunits(filehead, type; distance=1.0, mp=1.0, me=1.0) + setunits(head, type; distance=1.0, mp=1.0, me=1.0) Set the units for the output files. If type is given as "SI", "CGS", "NORMALIZED", "PIC", "PLANETARY", "SOLAR", set @@ -467,13 +480,13 @@ current density [jSI] in SI units. Distance unit [rplanet | rstar], ion and elec Also calculate convenient constants ti0, cs0 ... for typical formulas. This function is currently not used anywhere! """ -function setunits(filehead, type; distance=1.0, mp=1.0, me=1.0) - ndim = filehead.ndim - headline = filehead.headline - neqpar = filehead.neqpar - nw = filehead.nw - eqpar = filehead.eqpar - variables = filehead.variables +function setunits(head, type; distance=1.0, mp=1.0, me=1.0) + ndim = head.ndim + headline = head.headline + neqpar = head.neqpar + nw = head.nw + eqpar = head.eqpar + variables = head.variables mu0SI = 4π*1e-7 # H/m cSI = 2.9978e8 # speed of light, [m/s] @@ -660,7 +673,7 @@ function setunits(filehead, type; distance=1.0, mp=1.0, me=1.0) return true end -function Base.show(io::IO, data::BATLData) +function Base.show(io::IO, data::BATS) showhead(io, data) if data.list.bytes ≥ 1e9 println(io, "filesize: $(data.list.bytes/1e9) GB") @@ -675,7 +688,7 @@ function Base.show(io::IO, data::BATLData) end """ - showhead(file, filehead) + showhead(file, head) Displaying file header information. """ @@ -707,5 +720,5 @@ end Display file information of `data`. """ -showhead(data::BATLData) = showhead(data.list, data.head) -showhead(io, data::BATLData) = showhead(data.list, data.head, io) \ No newline at end of file +showhead(data::BATS) = showhead(data.list, data.head) +showhead(io, data::BATS) = showhead(data.list, data.head, io) \ No newline at end of file diff --git a/src/plot/plots.jl b/src/plot/plots.jl index 5a4fe5d1..3e8210dc 100644 --- a/src/plot/plots.jl +++ b/src/plot/plots.jl @@ -3,7 +3,7 @@ using RecipesBase # Build a recipe which acts on a custom type. -@recipe function f(bd::BATLData{1, T}, var::AbstractString) where {T} +@recipe function f(bd::BATS{1, 2, T}, var::AbstractString) where {T} hasunits = hasunit(bd) if hasunits @@ -22,7 +22,7 @@ using RecipesBase end end -@recipe function f(bd::BATLData{2, T}, var::AbstractString; +@recipe function f(bd::BATS{2, 3, T}, var::AbstractString; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1) where {T} hasunits = hasunit(bd) diff --git a/src/plot/pyplot.jl b/src/plot/pyplot.jl index 77074b86..a007c4e0 100644 --- a/src/plot/pyplot.jl +++ b/src/plot/pyplot.jl @@ -34,7 +34,7 @@ function plotlogdata(data, head::NamedTuple, func::AbstractString; plotmode="lin plotmode = split(plotmode) for (ivar, var) in enumerate(vars) - varIndex_ = findfirst(x->x==lowercase(var), lowercase.(head.variables)) + varIndex_ = findfirst(x->lowercase(x)==lowercase(var), head.variables) isnothing(varIndex_) && error("$(var) not found in file header variables!") figure() @@ -53,11 +53,11 @@ function plotlogdata(data, head::NamedTuple, func::AbstractString; plotmode="lin end """ - plotgrid(bd::BATLData, var, ax=nothing; kwargs...) + plotgrid(bd::BATS, var, ax=nothing; kwargs...) Plot 2D mesh. """ -function plotgrid(bd::BATLData{2, T}, func::AbstractString, ax=nothing; kwargs...) where T +function plotgrid(bd::BATS{2, 3, T}, func::AbstractString, ax=nothing; kwargs...) where T if isnothing(ax) ax = plt.gca() end # This does not take subdomain plot into account! @@ -81,7 +81,7 @@ end 2D plane cut pcolormesh of 3D box data. `sequence` is the index along `dir`. """ -function cutplot(bd::BATLData{3, T}, var::AbstractString, ax=nothing; +function cutplot(bd::BATS{3, 4, T}, var::AbstractString, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], dir="x", sequence=1) where {T} x, w = bd.x, bd.w varIndex_ = findindex(bd, var) @@ -132,13 +132,13 @@ end """ - streamslice(data::BATLData, var, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], dir="x", + streamslice(data::BATS, var, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], dir="x", sequence=1; kwargs...) Plot streamlines on 2D slices of 3D box data. Variable names in `var` string must be separated with `;`. """ -function streamslice(bd::BATLData{3, T}, var::AbstractString, ax=nothing; +function streamslice(bd::BATS{3, T}, var::AbstractString, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], dir="x", sequence=1, kwargs...) where {T} x,w = bd.x, bd.w varstream = split(var, ";") @@ -192,11 +192,11 @@ end """ - plot(bd::BATLData{1, T}, var, ax=nothing; kwargs...) + plot(bd::BATS{1, T}, var, ax=nothing; kwargs...) Wrapper over `plot` in matplotlib. Plot 1D outputs. """ -function PyPlot.plot(bd::BATLData{1, T}, var::AbstractString, ax=nothing; kwargs...) where T +function PyPlot.plot(bd::BATS{1, 2, T}, var::AbstractString, ax=nothing; kwargs...) where T x, w = bd.x, bd.w varIndex_ = findindex(bd, var) if isnothing(ax) ax = plt.gca() end @@ -219,7 +219,7 @@ end Wrapper over `scatter` in matplotlib. """ -function PyPlot.scatter(bd::BATLData{1, T}, var::AbstractString, ax=nothing; kwargs...) where T +function PyPlot.scatter(bd::BATS{1, 2, T}, var::AbstractString, ax=nothing; kwargs...) where T x, w = bd.x, bd.w varIndex_ = findindex(bd, var) if isnothing(ax) ax = plt.gca() end @@ -233,7 +233,7 @@ end Wrapper over `contour` in matplotlib. """ -function PyPlot.contour(bd::BATLData{2, T}, var::AbstractString, ax=nothing; levels=0, +function PyPlot.contour(bd::BATS{2, 3, T}, var::AbstractString, ax=nothing; levels=0, plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, rbody=1.0, kwargs...) where T Xi, Yi, Wi = interp2d(bd, var, plotrange, plotinterval; innermask, rbody) @@ -256,7 +256,7 @@ end Wrapper over `contourf` in matplotlib. See [`interp2d`](@ref) for some related keywords. """ -function PyPlot.contourf(bd::BATLData{2, T}, var::AbstractString, ax=nothing; levels::Int=0, +function PyPlot.contourf(bd::BATS{2, 3, T}, var::AbstractString, ax=nothing; levels::Int=0, plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, rbody=1.0, add_colorbar=true, vmin=-Inf, vmax=Inf, colorscale=:linear, kwargs...) where T Xi, Yi, Wi = interp2d(bd, var, plotrange, plotinterval; innermask, rbody) @@ -279,7 +279,7 @@ end Wrapper over `tricontourf` in matplotlib. """ -function PyPlot.tricontourf(bd::BATLData{2, T}, var::AbstractString, ax=nothing; +function PyPlot.tricontourf(bd::BATS{2, 3, T}, var::AbstractString, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], kwargs...) where T x, w = bd.x, bd.w varIndex_ = findindex(bd, var) @@ -305,7 +305,7 @@ function PyPlot.tricontourf(bd::BATLData{2, T}, var::AbstractString, ax=nothing; c end -function PyPlot.triplot(bd::BATLData{2, T}, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], +function PyPlot.triplot(bd::BATS{2, 3, T}, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], kwargs...) where T X = vec(bd.x[:,:,1]) Y = vec(bd.x[:,:,2]) @@ -323,12 +323,12 @@ function PyPlot.triplot(bd::BATLData{2, T}, ax=nothing; plotrange=[-Inf,Inf,-Inf end """ - plot_trisurf(data::BATLData, var::String, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], + plot_trisurf(data::BATS, var::String, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], kwargs...) Wrapper over `plot_trisurf` in matplotlib. """ -function PyPlot.plot_trisurf(bd::BATLData{2, T}, var::AbstractString; +function PyPlot.plot_trisurf(bd::BATS{2, 3, T}, var::AbstractString; plotrange=[-Inf,Inf,-Inf,Inf], kwargs...) where T x, w = bd.x, bd.w varIndex_ = findindex(bd, var) @@ -362,7 +362,7 @@ end Wrapper over `plot_surface` in matplotlib. """ -function PyPlot.plot_surface(bd::BATLData{2, T}, var::AbstractString, ax=nothing; +function PyPlot.plot_surface(bd::BATS{2, 3, T}, var::AbstractString, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, rbody=1.0, kwargs...) where T if isnothing(ax) ax = plt.gca() end @@ -389,7 +389,7 @@ end Wrapper over `pcolormesh` in matplotlib. """ -function PyPlot.pcolormesh(bd::BATLData{2, T}, var::AbstractString, ax=nothing; +function PyPlot.pcolormesh(bd::BATS{2, 3, T}, var::AbstractString, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, rbody=1.0, vmin=-Inf, vmax=Inf, colorscale=:linear, add_colorbar=true, kwargs...) where T xi, yi, Wi = interp2d(bd, var, plotrange, plotinterval; innermask, rbody) @@ -412,7 +412,7 @@ end Wrapper over `tripcolor` in matplotlib. """ -function PyPlot.tripcolor(bd::BATLData{2, T}, var::AbstractString, ax=nothing; +function PyPlot.tripcolor(bd::BATS{2, 3, T}, var::AbstractString, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], innermask=false, kwargs...) where {T} x, w = bd.x, bd.w @@ -461,7 +461,7 @@ end Wrapper over `streamplot` in matplotlib. """ -function PyPlot.streamplot(bd::BATLData{2, T}, var::AbstractString, ax=nothing; +function PyPlot.streamplot(bd::BATS{2, 3, T}, var::AbstractString, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=Inf, kwargs...) where T xi, yi, v1, v2 = _getvector(bd, var; plotrange, plotinterval) @@ -470,13 +470,12 @@ function PyPlot.streamplot(bd::BATLData{2, T}, var::AbstractString, ax=nothing; ax.streamplot(xi, yi, v1, v2; kwargs...) end -function _getvector(bd::BATLData{2, T}, var::AbstractString; +function _getvector(bd::BATS{2, 3, T}, var::AbstractString; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=Inf) where T x, w = bd.x, bd.w varstream = split(var, ";") - wnames = lowercase.(bd.head.wnames) - var1_ = findfirst(x->x==lowercase(varstream[1]), wnames) - var2_ = findfirst(x->x==lowercase(varstream[2]), wnames) + var1_ = findfirst(x->lowercase(x)==lowercase(varstream[1]), bd.head.wnames) + var2_ = findfirst(x->lowercase(x)==lowercase(varstream[2]), bd.head.wnames) if isinf(plotinterval) plotinterval = (x[end,1,1] - x[1,1,1]) / size(x, 1) end @@ -529,7 +528,7 @@ end Wrapper over `quiver` in matplotlib. Only supports Cartesian grid for now. """ -function PyPlot.quiver(bd::BATLData{2, T}, var::AbstractString, ax=nothing; +function PyPlot.quiver(bd::BATS{2, 3, T}, var::AbstractString, ax=nothing; stride::Integer=10, kwargs...) where T x, w = bd.x, bd.w VarQuiver = split(var, ";") @@ -567,7 +566,7 @@ function set_colorbar(colorscale, vmin, vmax, data=[1.0]) cnorm end -function add_titles(bd::BATLData, var, ax) +function add_titles(bd::BATS, var, ax) varIndex_ = findindex(bd, var) title(bd.head.wnames[varIndex_]) diff --git a/src/select.jl b/src/select.jl index b62755a5..a4ab4951 100644 --- a/src/select.jl +++ b/src/select.jl @@ -6,9 +6,9 @@ Get 2D plane cut in orientation `dir` for `var` out of 3D box `data` within `plotrange`. The returned 2D data lies in the `sequence` plane from - to + in `dir`. """ -function cutdata(bd::BATLData, var::AbstractString; +function cutdata(bd::BATS, var::AbstractString; plotrange=[-Inf,Inf,-Inf,Inf], dir::String="x", sequence::Int=1) - var_ = findfirst(x->x==lowercase(var), lowercase.(bd.head.wnames)) + var_ = findfirst(x->lowercase(x)==lowercase(var), bd.head.wnames) isempty(var_) && error("$(var) not found in header variables!") if dir == "x" @@ -163,7 +163,7 @@ function subvolume(x, y, z, u, v, w, limits) end """ - getvar(bd::BATLData, var::AbstractString) -> Array + getvar(bd::BATS, var::AbstractString) -> Array Return variable data from string `var`. This is also supported via direct indexing, @@ -172,11 +172,11 @@ Return variable data from string `var`. This is also supported via direct indexi bd["rho"] ``` """ -function getvar(bd::BATLData{ndim, T}, var::AbstractString) where {ndim, T} +function getvar(bd::BATS{ndim, ndimp1, T}, var::AbstractString) where {ndim, ndimp1, T} if var in keys(variables_predefined) w = variables_predefined[var](bd) else - var_ = findfirst(x->x==lowercase(var), lowercase.(bd.head.wnames)) + var_ = findfirst(x->lowercase(x)==lowercase(var), bd.head.wnames) isnothing(var_) && error("$var not found in file header variables!") w = selectdim(bd.w, ndim+1, var_) end @@ -184,15 +184,14 @@ function getvar(bd::BATLData{ndim, T}, var::AbstractString) where {ndim, T} w end -@inline @Base.propagate_inbounds Base.getindex(bd::BATLData, var::AbstractString) = +@inline @Base.propagate_inbounds Base.getindex(bd::BATS, var::AbstractString) = getvar(bd, var) - "Construct vectors from scalar components." -function _fill_vector_from_scalars(bd::BATLData{ndim, T, U}, var) where {ndim, T, U} +function _fill_vector_from_scalars(bd::BATS{ndim, ndimp1, T}, var) where {ndim, ndimp1, T} v1, v2, v3 = _get_vectors(bd, var) - v = Array{T, ndims(v1)+1}(undef, 3, size(v1)...) + v = Array{T, ndimp1}(undef, 3, size(v1)...) Rpost = CartesianIndices(size(v1)) for Ipost in Rpost @@ -204,64 +203,58 @@ function _fill_vector_from_scalars(bd::BATLData{ndim, T, U}, var) where {ndim, T v end -function _get_magnitude2(bd::BATLData{2, T, U}, var=:B) where {T, U} +function _get_magnitude2(bd::BATS{2, 3, T}, var=:B) where T vx, vy, vz = _get_vectors(bd, var) v = similar(vx)::Array{T, 2} - @simd for i in eachindex(v) + @inbounds @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} +function _get_magnitude(bd::BATS{2, 3, T}, var=:B) where T vx, vy, vz = _get_vectors(bd, var) v = similar(vx)::Array{T, 2} - @simd for i in eachindex(v) + @inbounds @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) +function _get_vectors(bd::BATS{Dim, Dimp1, T}, var) where {Dim, Dimp1,T} + VT = SubArray{T, Dim, Array{T, Dimp1}, Tuple{Base.Slice{Base.OneTo{Int64}}, + Base.Slice{Base.OneTo{Int64}}, Int64}, true} if var == :B - vx = bd["Bx"] - vy = bd["By"] - vz = bd["Bz"] + vx, vy, vz = bd["Bx"]::VT, bd["By"]::VT, bd["Bz"]::VT elseif var == :E - vx = bd["Ex"] - vy = bd["Ey"] - vz = bd["Ez"] + vx, vy, vz = bd["Ex"]::VT, bd["Ey"]::VT, bd["Ez"]::VT elseif var == :U - vx = bd["Ux"] - vy = bd["Uy"] - vz = bd["Uz"] + vx, vy, vz = bd["Ux"]::VT, bd["Uy"]::VT, bd["Uz"]::VT elseif var == :U0 - vx = bd["UxS0"] - vy = bd["UyS0"] - vz = bd["UzS0"] + vx, vy, vz = bd["UxS0"]::VT, bd["UyS0"]::VT, bd["UzS0"]::VT elseif var == :U1 - vx = bd["UxS1"] - vy = bd["UyS1"] - vz = bd["UzS1"] + vx, vy, vz = bd["UxS1"]::VT, bd["UyS1"]::VT, bd["UzS1"]::VT 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"] +function _get_anisotropy(bd::BATS{2, 3, T}, species=0) where T + VT = SubArray{T, 2, Array{T, 3}, Tuple{Base.Slice{Base.OneTo{Int64}}, + Base.Slice{Base.OneTo{Int64}}, Int64}, true} + Bx, By, Bz = bd["Bx"]::VT, bd["By"]::VT, bd["Bz"]::VT # Rotate the pressure tensor to align the 3rd direction with B pop = string(species) - Pxx = bd["pXXS"*pop] - Pyy = bd["pYYS"*pop] - Pzz = bd["pZZS"*pop] - Pxy = bd["pXYS"*pop] - Pxz = bd["pXZS"*pop] - Pyz = bd["pYZS"*pop] + Pxx = bd["pXXS"*pop]::VT + Pyy = bd["pYYS"*pop]::VT + Pzz = bd["pZZS"*pop]::VT + Pxy = bd["pXYS"*pop]::VT + Pxz = bd["pXZS"*pop]::VT + Pyz = bd["pYZS"*pop]::VT Paniso = similar(Pxx)::Array{T, 2} @@ -289,6 +282,8 @@ const variables_predefined = Dict{String, Function}( "Bmag" => (bd -> _get_magnitude(bd, :B)), "Emag" => (bd -> _get_magnitude(bd, :E)), "Umag" => (bd -> _get_magnitude(bd, :U)), + "Uemag" => (bd -> _get_magnitude(bd, :U0)), + "Uimag" => (bd -> _get_magnitude(bd, :U1)), "B" => (bd -> _fill_vector_from_scalars(bd, :B)), "E" => (bd -> _fill_vector_from_scalars(bd, :E)), "U" => (bd -> _fill_vector_from_scalars(bd, :U)), diff --git a/src/type.jl b/src/type.jl new file mode 100644 index 00000000..8b0c4af6 --- /dev/null +++ b/src/type.jl @@ -0,0 +1,49 @@ +abstract type AbstractBATS{Dim, T} end + +@enum FileType Real4Bat=1 Real8Bat=2 AsciiBat=3 LogBat=4 TecBat=5 + +"Type for Batsrus file information." +struct FileList + "filename" + name::String + "file type" + type::FileType + "directory" + dir::String + "file size" + bytes::Int + "number of snapshots" + npictinfiles::Int + "length of meta data" + lenhead::Int +end + +"Batsrus file head information." +struct BatsHead + ndim::Int + headline::SubString{String} + it::Int + time::Float32 + gencoord::Bool + neqpar::Int + nw::Int + nx::Vector{Int} + eqpar::Vector{Float32} + variables::Vector{SubString{String}} + wnames::Vector{SubString{String}} +end + +"Batsrus data container. `Dim` is the dimension of output, and `Dimp1` is an extra parameter for working around derived types." +struct BATS{Dim, Dimp1, T<:AbstractFloat} <: AbstractBATS{Dim, T} + head::BatsHead + list::FileList + "Grid" + x::Array{T, Dimp1} + "Variables" + w::Array{T, Dimp1} + + function BATS(head, list, x::Array{T, Dimp1}, w::Array{T, Dimp1}) where {T, Dimp1} + @assert head.ndim + 1 == Dimp1 "Dimension mismatch!" + new{Dimp1-1, Dimp1, T}(head, list, x, w) + end +end \ No newline at end of file diff --git a/src/unit/UnitfulBatsrus.jl b/src/unit/UnitfulBatsrus.jl index c976dfe5..095f2c51 100644 --- a/src/unit/UnitfulBatsrus.jl +++ b/src/unit/UnitfulBatsrus.jl @@ -48,7 +48,7 @@ end function getunit(data, var) # Batsrus has a bug in the 2D cuts of 3D runs: it always outputs the 3 # coordinate units in the headline. To work around it, here the index is shifted by 1. - var_ = findfirst(x->x==lowercase(var), lowercase.(data.head.variables)) + 1 + var_ = findfirst(x->lowercase(x)==lowercase(var), data.head.variables) + 1 isnothing(var_) && error("$(var) not found in file header variables!") if data.head.headline in ("normalized variables", "PLANETARY") var_unit = nothing diff --git a/src/utility.jl b/src/utility.jl index 680ac852..1419d09d 100644 --- a/src/utility.jl +++ b/src/utility.jl @@ -1,7 +1,7 @@ # Utility functions for plotting and analyzing. """ - interp2d(bd::BATLData, var::AbstractString, plotrange=[-Inf, Inf, -Inf, Inf], + interp2d(bd::BATS, var::AbstractString, plotrange=[-Inf, Inf, -Inf, Inf], plotinterval=Inf; kwargs...) Return 2D interpolated slices of data `var` from `bd`. If `plotrange` is not set, output @@ -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, U}, var::AbstractString, +function interp2d(bd::BATS{2, 3, T}, var::AbstractString, plotrange::Vector=[-Inf, Inf, -Inf, Inf], plotinterval::Real=Inf; - innermask::Bool=false, rbody::Real=1.0, useMatplotlib::Bool=true) where {T, U} + innermask::Bool=false, rbody::Real=1.0, useMatplotlib::Bool=true) where T x, w = bd.x, bd.w varIndex_ = findindex(bd, var) @@ -87,8 +87,8 @@ function interp2d(bd::BATLData{2, T, U}, var::AbstractString, end "Find variable index in the BATSRUS data." -function findindex(bd::BATLData, var::AbstractString) - varIndex_ = findfirst(x->x==lowercase(var), lowercase.(bd.head.wnames)) +function findindex(bd::BATS, var::AbstractString) + varIndex_ = findfirst(x->lowercase(x)==lowercase(var), bd.head.wnames) isnothing(varIndex_) && error("$(var) not found in file header variables!") varIndex_ @@ -102,7 +102,7 @@ function meshgrid(x, y) X, Y end -@inline function hasunit(bd::BATLData) +@inline function hasunit(bd::BATS) startswith(bd.head.headline, "normalized") ? false : true end @@ -129,11 +129,11 @@ function interpolate2d_generalized_coords(X::T, Y::T, W::T, end """ - interp1d(bd::BATLData, var::AbstractString, loc::AbstractVector{<:AbstractFloat}) + interp1d(bd::BATS, var::AbstractString, loc::AbstractVector{<:AbstractFloat}) Interpolate `var` at spatial point `loc` in `bd`. """ -function interp1d(bd::BATLData{2, T, U}, var::AbstractString, loc::AbstractVector{<:AbstractFloat}) where {T, U} +function interp1d(bd::BATS{2, 3, T}, var::AbstractString, loc::AbstractVector{<:AbstractFloat}) where T @assert !bd.head.gencoord "Only accept structured grids!" x = bd.x @@ -146,11 +146,11 @@ function interp1d(bd::BATLData{2, T, U}, var::AbstractString, loc::AbstractVecto end """ - interp1d(bd::BATLData, var::AbstractString, point1::Vector, point2::Vecto) + interp1d(bd::BATS, var::AbstractString, point1::Vector, point2::Vecto) Interpolate `var` along a line from `point1` to `point2` in `bd`. """ -function interp1d(bd::BATLData{2, T, U}, var::AbstractString, point1::Vector, point2::Vector) where {T, U} +function interp1d(bd::BATS{2, 3, T}, var::AbstractString, point1::Vector, point2::Vector) where T @assert !bd.head.gencoord "Only accept structured grids!" x = bd.x @@ -179,20 +179,20 @@ function slice1d(bd, var, icut::Int=1, dir::Int=2) end "Return view of variable `var` in `bd`." -function getview(bd::BATLData{1, T, U}, var) where {T, U} +function getview(bd::BATS{1, 2, T}, var) where T varIndex_ = findindex(bd, var) v = @view bd.w[:,varIndex_] end -function getview(bd::BATLData{2, T, U}, var) where {T, U} +function getview(bd::BATS{2, 3, T}, var) where T varIndex_ = findindex(bd, var) v = @view bd.w[:,:,varIndex_] end "Return value range of `var` in `bd`." -get_var_range(bd::BATLData, var) = getview(bd, var) |> extrema +get_var_range(bd::BATS, var) = getview(bd, var) |> extrema "Squeeze singleton dimensions for an array `A`." function squeeze(A::AbstractArray) @@ -201,7 +201,7 @@ function squeeze(A::AbstractArray) dropdims(A, dims=singleton_dims) end -function get_range(bd::BATLData, var::AbstractString, verbose=true) +function get_range(bd::BATS, var::AbstractString, verbose=true) varIndex_ = findindex(bd, var) w = selectdim(bd.w, bd.head.ndim+1, varIndex_) wmin, wmax = extrema(w) diff --git a/test/runtests.jl b/test/runtests.jl index e3efba57..3c9d720e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -44,6 +44,7 @@ end @testset "Reading 2D structured binary" begin file = "z=0_raw_1_t25.60000_n00000258.out" + @test_throws ArgumentError load(joinpath(datapath, file), npict=2) bd = load(joinpath(datapath, file)) @test bd.head.time == 25.6f0 @test extrema(bd.x) == (-127.5f0, 127.5f0) @@ -53,6 +54,7 @@ end @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 + @test bd["B2"][128,2] == 0.8507747f0 # Linear interpolation at a given point d = interp1d(bd, "rho", Float32[0.0, 0.0]) @test d == 0.6936918f0 @@ -67,17 +69,23 @@ end file = "z=0_fluid_region0_0_t00001640_n00010142.out" bd = load(joinpath(datapath, file)) + @test bd["Emag"][2,1] == 2655.4805f0 + @test bd["E2"][2,1] == 7.051577f6 + @test bd["E"][:,2,1] == Float32[-241.05942, -2644.2058, -40.53219] + @test bd["Ue2"][2,1] == 33784.973f0 + @test bd["Ui2"][2,1] == bd["Ui2"][2,1] @test bd["Anisotropy0"][1:2,1] ≈ Float32[1.2630985, 2.4700143] @test bd["Anisotropy1"][1:2,1] ≈ Float32[1.2906302, 2.6070855] end - @testset "Reading 2D unstructured" begin + @testset "Reading 2D unstructured ascii" begin file = "bx0_mhd_6_t00000100_n00000352.out" bd = load(joinpath(datapath, file)) 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 + @test bd["U2"][2] == 5163.073119959886 end @testset "Reading 3D structured binary" begin