Separate raw data extraction and derived data extraction

henry2004y · Dec 15, 2024 · 7b6a076 · 7b6a076
1 parent 6c7d8ef
commit 7b6a076
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 65 deletions.
diff --git a/docs/src/man/manual.md b/docs/src/man/manual.md
@@ -57,26 +57,24 @@ w = interp1d(bd, "rho", point1, point2)
 
 - Derived variables
 
+We provide utility methods `get_magnitude`, `get_magnitude2`, and `fill_vector_from_scalars` for vector processing:
+
 ```julia
-Bmag = bd["Bmag"]
+Bmag = get_magnitude(bd, :B)
+B2 = get_magnitude2(bd, :B)
+Bvec = Batsrus.fill_vector_from_scalars(bd, :B)
+paniso0 = get_anisotropy(bd, 0)
 ```
 
-Here is a full list of predefined derived quantities:
+These are built upon `get_vectors`. Here is a full list of predefined derived quantities in `get_vectors`:
 
 | 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        |
-| 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[0-1]       | Pressure/Temperature anisotropy  | B, P tensor       |
+| :B                    | Magnetic field vector            | Bx, By, Bz        |
+| :E                    | Electric field vector            | Ex, Ey, Ez        |
+| :U                    | Velocity vector                  | Ux, Uy, Uz        |
+| :U0                   | Electron velocity vector         | UxS0, UyS0, UzS0  |
+| :U1                   | Proton velocity vector           | UxS1, UyS1, UzS1  |
 
 ## Output format conversion
 

diff --git a/src/Batsrus.jl b/src/Batsrus.jl
@@ -13,7 +13,7 @@ using DimensionalData
 
 export BATS,
    load, readlogdata, readtecdata, showhead, # io
-   getvar, cutdata, subvolume, subsurface, get_convection_E, get_hall_E, 
+   getvar, getderivedvar, cutdata, subvolume, subsurface, get_convection_E, get_hall_E, 
    get_anisotropy, get_vectors, get_magnitude, get_magnitude2,
    fill_vector_from_scalars, # select
    Batl, convertTECtoVTU, convertIDLtoVTK, readhead, readtree, getConnectivity, # vtk

diff --git a/src/select.jl b/src/select.jl
@@ -173,41 +173,7 @@ bd["rho"]
 ```
 """
 function getvar(bd::BATS{ndim, TV, TX, TW}, var::AbstractString) where {ndim, TV, TX, TW}
-   if var == "B2"
-      w = get_magnitude2(bd, :B)
-   elseif var == "E2"
-      w = get_magnitude2(bd, :E)
-   elseif var == "U2"
-      w = get_magnitude2(bd, :U)
-   elseif var == "Ue2"
-      w = get_magnitude2(bd, :U0)
-   elseif var == "Ui2"
-      w = get_magnitude2(bd, :U1)
-   elseif var == "Bmag"
-      w = get_magnitude(bd, :B)
-   elseif var == "Emag"
-      w = get_magnitude(bd, :E)
-   elseif var == "Umag"
-      w = get_magnitude(bd, :U)
-   elseif var == "Uemag"
-      w = get_magnitude(bd, :U0)
-   elseif var == "Uimag"
-      w = get_magnitude(bd, :U1)
-   elseif var == "B"
-      w = fill_vector_from_scalars(bd, :B)
-   elseif var == "E"
-      w = fill_vector_from_scalars(bd, :E)
-   elseif var == "U"
-      w = fill_vector_from_scalars(bd, :U)
-   elseif var == "Anisotropy0"
-      w = get_anisotropy(bd, 0)
-   elseif var == "Anisotropy1"
-      w = get_anisotropy(bd, 1)
-   else
-      w = @view bd.w[var=At(lowercase(var))]
-   end
-
-   w
+   w = @view bd.w[var=At(lowercase(var))]
 end
 
 @inline @Base.propagate_inbounds Base.getindex(bd::BATS, var::AbstractString) =
@@ -233,7 +199,7 @@ Calculate the magnitude square of vector `var`. See [`get_vectors`](@ref) for th
 """
 function get_magnitude2(bd::BATS, var=:B)
    vx, vy, vz = get_vectors(bd, var)
-   v = similar(vx)
+   v = similar(vx.data)
 
    @inbounds @simd for i in eachindex(v)
       v[i] = vx[i]^2 + vy[i]^2 + vz[i]^2
@@ -249,7 +215,7 @@ Calculate the magnitude of vector `var`. See [`get_vectors`](@ref) for the optio
 """
 function get_magnitude(bd::BATS, var=:B)
    vx, vy, vz = get_vectors(bd, var)
-   v = similar(vx)
+   v = similar(vx.data)
 
    @inbounds @simd for i in eachindex(v)
       v[i] = √(vx[i]^2 + vy[i]^2 + vz[i]^2)
@@ -295,7 +261,7 @@ function get_anisotropy(bd::BATS{2, TV, TX, TW}, species=0) where {TV, TX, TW}
    Pxz = bd["pXZS"*pop]
    Pyz = bd["pYZS"*pop]
    #TODO: Generalize to n-dimension with CartesianIndices
-   Paniso = similar(Pxx)
+   Paniso = similar(Pxx.data)
 
    @inbounds for j in axes(Pxx, 2), i in axes(Pxx, 1)  
       b̂ = normalize(SA[Bx[i,j], By[i,j], Bz[i,j]])

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -52,9 +52,9 @@ end
       plotrange = [-10.0, 10.0, -Inf, Inf]
       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
-      @test bd["B2"][128,2] == 0.8507747f0
+      @test Batsrus.fill_vector_from_scalars(bd, :B)[:,end,end] == Float32[1.118034, -0.559017, 0.0]
+      @test get_magnitude(bd, :B)[128,2] == 0.9223745f0
+      @test get_magnitude2(bd, :B)[128,2] == 0.8507747f0
       # Linear interpolation at a given point
       d = interp1d(bd, "rho", Float32[0.0, 0.0])
       @test d == 0.6936918f0
@@ -73,13 +73,12 @@ end
       @test length(x) == 601 && y[2] == 0.0f0
       x, y = Batsrus.meshgrid(bd, Float32[-100, 100, -Inf, Inf])
       @test length(x) == 4
-      @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]
+      @test get_magnitude(bd, :E)[2,1] == 2655.4805f0
+      @test get_magnitude2(bd, :E)[2,1] == 7.051577f6
+      @test Batsrus.fill_vector_from_scalars(bd, :E)[:,2,1] == Float32[-241.05942, -2644.2058, -40.53219]
+      @test get_magnitude2(bd, :U0)[2,1] == 33784.973f0
+      @test get_anisotropy(bd, 0)[1:2,1] ≈ Float32[1.2630985, 2.4700143]
+      @test get_anisotropy(bd, 1)[1:2,1] ≈ Float32[1.2906302, 2.6070855]
       w = get_convection_E(bd)
       @test w[2][2,1] ≈ -2454.3933f0
       w = get_hall_E(bd)
@@ -94,8 +93,8 @@ end
       @test length(x) == 117 && length(y) == 246
       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
+      @test get_magnitude(bd, :U)[2] == 71.85452748407637
+      @test get_magnitude2(bd, :U)[2] == 5163.073119959886
    end
 
    @testset "Reading 3D structured binary" begin