Lazy struct (#148)

* Refactor readvcells
henry2004y · Jul 17, 2023 · 4f9e488 · 4f9e488
1 parent 67948f2
commit 4f9e488
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 29 deletions.
diff --git a/Project.toml b/Project.toml
@@ -6,6 +6,7 @@ version = "0.10.2"
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
+LazilyInitializedFields = "0e77f7df-68c5-4e49-93ce-4cd80f5598bf"
 LazyGrids = "7031d0ef-c40d-4431-b2f8-61a8d2f650db"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Mmap = "a63ad114-7e13-5084-954f-fe012c677804"
@@ -20,6 +21,7 @@ XML = "72c71f33-b9b6-44de-8c94-c961784809e2"
 
 [compat]
 LaTeXStrings = "1.2"
+LazilyInitializedFields = "1.0"
 LazyGrids = "0.4, 0.5"
 Parsers = "2.4"
 PrecompileTools = "1.0"

diff --git a/src/Vlasiator.jl b/src/Vlasiator.jl
@@ -7,6 +7,7 @@ using Statistics: mean
 using XML: Node, LazyNode, Element, Document, children, tag, attributes, value
 import XML
 using Mmap: mmap
+using LazilyInitializedFields
 using WriteVTK
 using LazyGrids: ndgrid
 using LaTeXStrings

diff --git a/src/vlsv/vlsvreader.jl b/src/vlsv/vlsvreader.jl
@@ -3,14 +3,15 @@
 const NodeVector = SubArray{Node, 1, Vector{Node}, Tuple{UnitRange{Int64}},
    true}
 
-"Velocity mesh information."
-struct VMeshInfo
+@lazy struct VMeshInfo
    "number of velocity blocks"
    vblocks::NTuple{3, Int}
    vblock_size::NTuple{3, Int}
    vmin::NTuple{3, Float64}
    vmax::NTuple{3, Float64}
    dv::NTuple{3, Float64}
+   @lazy cellsWithVDF::Vector{Int}
+   @lazy nblock_C::Vector{Int}
 end
 
 "Variable information from the VLSV footer."
@@ -307,8 +308,8 @@ function load(file::AbstractString)
          vmin[1], vmin[2], vmin[3] = nodeX[begin], nodeY[begin], nodeZ[begin]
          vmax[1], vmax[2], vmax[3] = nodeX[end], nodeY[end], nodeZ[end]
       else
-         popname = "avgs"
          # In VLSV before 5.0 the mesh is defined with parameters.
+         popname = "avgs"
          if "vxblocks_ini" in getindex.(n.param, "name")
             vblocks_str = ("vxblocks_ini", "vyblocks_ini", "vzblocks_ini")
             vmin_str = ("vxmin", "vymin", "vzmin")
@@ -328,9 +329,9 @@ function load(file::AbstractString)
          push!(species, popname)
       end
 
-      # Create a new object for this population
+      # Create a new object for this species
       popVMesh = VMeshInfo(NTuple{3}(vblocks), NTuple{3}(vblock_size),
-         NTuple{3}(vmin), NTuple{3}(vmax), dv)
+         NTuple{3}(vmin), NTuple{3}(vmax), dv, uninit, uninit)
 
       meshes[popname] = popVMesh
    end
@@ -825,9 +826,7 @@ function readvcells(meta::MetaVLSV, cid::Int; species::String="proton")
    (;vblock_size) = meta.meshes[species]
    bsize = prod(vblock_size)
 
-   offset_v, nblocks = 0, 0
-
-   let cellsWithVDF, nblock_C
+   if !@isinit meta.meshes[species].cellsWithVDF
       for node in nodeVLSV.cellwithVDF
          at = attributes(node)
          if at["name"] == species
@@ -836,6 +835,7 @@ function readvcells(meta::MetaVLSV, cid::Int; species::String="proton")
             cellsWithVDF = Vector{Int}(undef, asize)
             seek(fid, offset)
             read!(fid, cellsWithVDF)
+            @init! meta.meshes[species].cellsWithVDF = cellsWithVDF
             break
          end
       end
@@ -846,48 +846,48 @@ function readvcells(meta::MetaVLSV, cid::Int; species::String="proton")
             asize = Parsers.parse(Int, at["arraysize"])
             dsize = Parsers.parse(Int, at["datasize"])
             offset = Parsers.parse(Int, value(node[1]))
-            nblock_C = dsize == 4 ?
-               Vector{Int32}(undef, asize) : Vector{Int}(undef, asize)
+            T = dsize == 4 ? Int32 : Int
+            nblock_C = Vector{Int}(undef, asize)
             seek(fid, offset)
-            read!(fid, nblock_C)
+            @inbounds for i in 1:asize
+               nblock_C[i] = read(fid, T)
+            end
+            @init! meta.meshes[species].nblock_C = nblock_C
             break
          end
       end
+   end
 
-      cellWithVDFIndex = findfirst(==(cid), cellsWithVDF)
-      if !isnothing(cellWithVDFIndex)
-         @inbounds nblocks = Int(nblock_C[cellWithVDFIndex])
-         if nblocks == 0
-            throw(ArgumentError("Cell ID $cid does not store velocity distribution!"))
-         end
-      else
+   cellWithVDFIndex = findfirst(==(cid), meta.meshes[species].cellsWithVDF)
+   if !isnothing(cellWithVDFIndex)
+      @inbounds nblocks = meta.meshes[species].nblock_C[cellWithVDFIndex]
+      if nblocks == 0
          throw(ArgumentError("Cell ID $cid does not store velocity distribution!"))
       end
-      # Offset position to vcell storage
-      @inbounds for i in 1:cellWithVDFIndex-1
-         offset_v += nblock_C[i]
-      end
+   else
+      throw(ArgumentError("Cell ID $cid does not store velocity distribution!"))
    end
+   # Offset position to vcell storage
+   offset_v = sum(@view meta.meshes[species].nblock_C[1:cellWithVDFIndex-1]; init=0)
 
-   local dsize, vsize, offset
-   # Read in avgs
+   local dsize, offset
+   # Read raw VDF
    for node in nodeVLSV.blockvar
       at = attributes(node)
       if at["name"] == species
          dsize = Parsers.parse(Int, at["datasize"])
-         vsize = Parsers.parse(Int, at["vectorsize"])
          offset = Parsers.parse(Int, value(node[1]))
          break
       end
    end
 
    data = let
-      Tavg = dsize == 4 ? Float32 : Float64
-      a = mmap(fid, Vector{UInt8}, dsize*vsize*nblocks, offset_v*vsize*dsize + offset)
-      reshape(reinterpret(Tavg, a), vsize, nblocks)
+      T = dsize == 4 ? Float32 : Float64
+      a = mmap(fid, Vector{UInt8}, dsize*bsize*nblocks, offset_v*bsize*dsize + offset)
+      reshape(reinterpret(T, a), bsize, nblocks)
    end
 
-   # Read in block IDs
+   # Read block IDs
    for node in nodeVLSV.blockid
       at = attributes(node)
       if at["name"] == species