Add predefined variables (#43)

Project.toml

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

docs/make.jl

@@ -12,7 +12,7 @@ makedocs(;
     ),
     pages=[
         "Home" => "index.md",
-        "Example" => "man/examples.md",
+        "Manual" => "man/manual.md",
         "Internal" => "man/internal.md",
         "Log" => "man/log.md",
     ],

docs/src/man/examples.md → docs/src/man/manual.md

@@ -38,9 +38,23 @@ bd["rho"]
 
 ```julia
 v = getvars(bd, ["Bx", "By", "Bz"])
-B = @. sqrt(v["Bx"]^2 + v["By"]^2 + v["Bz"]^2)
+Bmag = bd["Bmag"]
 ```
 
+Here is a full list of predefined derived quantities:
+
+| Derived variable name | Meaning                          | Required variable |
+|-----------------------|----------------------------------|-------------------|
+| B2                    | magnetic field magnitude squared | Bx, By, Bz        |
+| E2                    | electric field magnitude squared | Ex, Ey, Ez        |
+| U2                    | velocity magnitude squared       | Ux, Uy, Uz        |
+| Bmag                  | magnetic field magnitude         | Bx, By, Bz        |
+| Emag                  | electric field magnitude         | Ex, Ey, Ez        |
+| Umag                  | velocity magnitude               | Ux, Uy, Uz        |
+| B                     | magnetic field vector            | Bx, By, Bz        |
+| E                     | electric field vector            | Ex, Ey, Ez        |
+| U                     | velocity vector                  | Ux, Uy, Uz        |
+
 ## Output format conversion
 
 We can convert 2D/3D BATSRUS outputs `*.dat` to VTK formats. It uses the VTK XML format writer [writeVTK](https://github.com/jipolanco/WriteVTK.jl) to generate files for Paraview and Tecplot. The default converted filename is `out.vtu`.
@@ -49,8 +63,6 @@ ASCII Tecplot file (supports both `tec` and `tcp`) and binary Tecplot file (set
 
 ```julia
 file = "x=0_mhd_1_n00000050.dat"
-#file = "3d_ascii.dat"
-#file = "3d_bin.dat"
 head, data, connectivity = readtecdata(file)
 convertTECtoVTU(head, data, connectivity)
 ```

src/Batsrus.jl

@@ -39,6 +39,9 @@ struct BATLData{T<:AbstractFloat}
    list::FileList
 end
 
+include("unit/UnitfulBatsrus.jl")
+using .UnitfulBatsrus
+
 include("io.jl")
 include("select.jl")
 include("vtk.jl")
@@ -48,9 +51,6 @@ include("plot/plots.jl")
 include("hdf.jl")
 @reexport using .HDF
 
-include("unit/UnitfulBatsrus.jl")
-using .UnitfulBatsrus
-
 function __init__()
    @require PyPlot="d330b81b-6aea-500a-939a-2ce795aea3ee" begin
       include("plot/pyplot.jl")

src/select.jl

@@ -177,10 +177,15 @@ bd["rho"]
 See also: [`getvars`](@ref).
 """
 function getvar(bd::BATLData, var::AbstractString)
-   var_ = findfirst(x->x==lowercase(var), lowercase.(bd.head.wnames))
-   isnothing(var_) && error("$var not found in file header variables!")
+   if var in keys(variables_predefined)
+      w = variables_predefined[var](bd)
+   else
+      var_ = findfirst(x->x==lowercase(var), lowercase.(bd.head.wnames))
+      isnothing(var_) && error("$var not found in file header variables!")
+      w = selectdim(bd.w, bd.head.ndim+1, var_)
+   end
 
-   w = selectdim(bd.w, bd.head.ndim+1, var_)
+   w
 end
 
 @inline @Base.propagate_inbounds Base.getindex(bd::BATLData, var::AbstractString) =
@@ -201,10 +206,34 @@ function getvars(bd::BATLData{U}, Names::Vector{T}) where {U, T<:AbstractString}
    dict
 end
 
+"Construct vectors from scalar components."
+function _fill_vector_from_scalars(bd::BATLData, vstr1, vstr2, vstr3)
+   v1 = getvar(bd, vstr1)
+   v2 = getvar(bd, vstr2)
+   v3 = getvar(bd, vstr3)
+   v = Array{eltype(v1), ndims(v1)+1}(undef, 3, size(v1)...)
+
+   Rpost = CartesianIndices(size(v1))
+   for Ipost in Rpost
+      v[1,Ipost] = v1[Ipost]
+      v[2,Ipost] = v2[Ipost]
+      v[3,Ipost] = v3[Ipost]
+   end
+
+   v
+end
 
+# Define derived parameters
 const variables_predefined = Dict(
-   "B" => bd -> sqrt.(getvar(bd, "Bx").^2 .+ getvar(bd, "By").^2 .+ getvar(bd, "Bz").^2),
-   "E" => bd -> sqrt.(getvar(bd, "Ex").^2 .+ getvar(bd, "Ey").^2 .+ getvar(bd, "Ez").^2),
-   "U" => bd -> sqrt.(getvar(bd, "Ux").^2 .+ getvar(bd, "Uy").^2 .+ getvar(bd, "Uz").^2),
-   #"beta" => bd -> getvar(bd, "P") ./ getvar(bd, "B").^2 * 2μ,
+   "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")),
 )

test/runtests.jl

@@ -52,6 +52,7 @@ end
       plotrange = [-10.0, 10.0, -Inf, Inf]
       x, y, w = Batsrus.getdata2d(bd, "rho", plotrange)
       @test w[1,end] == 0.6848978549242021
+      @test bd["B"][:,end,end] == Float32[1.118034, -0.559017, 0.0]
    end
 
    @testset "Reading 2D unstructured binary" begin