Preparing new release

Project.toml

@@ -19,11 +19,11 @@ Accessors = "0.1"
 Corpuscles = "2"
 FHist = "0.10"
 Graphs = "1"
+PrettyTables = "2"
 StaticArrays = "1.9"
 StructArrays = "0.6"
 UnROOT = "0.10"
 YAML = "0.4"
-PrettyTables = "2"
 julia = "1"
 
 [extras]

docs/src/release_notes.md

@@ -1,6 +1,13 @@
 
 # Release Notes
 
+
+## 0.4.0
+### New Functionality
+- Added function `RootIO.create_getter(reader::Reader, bname::String; selection=nothing)` to create a getter function for a specific branch.
+  The optional argument allows to select leaves to be read.
+- The overall performance is highly improved (factor 3 with respect previous version)
+
 ## 0.3.1
 ### Bug Fixes
 - Legacy podio test fixed

examples/FCC/analysis_MT.jl

@@ -42,18 +42,16 @@ mutable struct MyData <: AbstractAnalysisData
 end
 
 get_μIDs  = RootIO.create_getter(reader, "Muon#0")
-get_recps = RootIO.create_getter(reader, "ReconstructedParticles")
+get_recps = RootIO.create_getter(reader, "ReconstructedParticles"; selection=[:energy,:momentum,:charge,:mass])
 
 function myanalysis!(data::MyData, reader, events)
     for evt in events
         data.pevts += 1                               # count process events
         #μIDs = RootIO.get(reader, evt, "Muon#0"; register=false) # get the ids of muons
-        #μIDs = RootIO.getCollection{ObjectID{ReconstructedParticle},Symbol("Muon#0")}(evt,UInt32(0)) # get the ids of muons
         μIDs = get_μIDs(evt)        
         length(μIDs) < 2 && continue                  # skip if less than 2  
 
         #recps = RootIO.get(reader, evt, "ReconstructedParticles"; register=false) 
-        #recps = RootIO.getCollection{ReconstructedParticle, :ReconstructedParticles}(evt, UInt32(15)) # get the reco particles
         recps = get_recps(evt)
         muons = recps[μIDs]                           # use the ids to subset the reco particles
 

podio/generate.jl

@@ -226,6 +226,7 @@ function gen_docstring(io, key, dtype)
     println(io,"\"\"\"")
 end
 
+#=
 function gen_structarray(io, key, dtype; podio=17)
     jtype = to_julia(key)
     println(io, "function StructArray{$jtype, bname}(evt::UnROOT.LazyEvent, collid = UInt32(0), len = -1) where bname")
@@ -326,6 +327,7 @@ function gen_structarray_rntuple(io, key, dtype)
     println(io, "    return StructArray{$jtype}(columns)")
     println(io, "end\n")
 end
+=#
 
 function build_graph(datatypes)
     types = to_julia.(keys(datatypes))
@@ -341,7 +343,6 @@ function build_graph(datatypes)
     graph
 end
 
-
 #---Components-------------------------------------------------------------------------------------
 io = open(joinpath(@__DIR__, "genComponents.jl"), "w")
 components = data["components"]
@@ -367,6 +368,7 @@ end
 println(io, "export $(join(unique(exports),", "))")
 close(io)
 
+#=
 #---StructArrays--------------------------------------------------------------------------------------
 for v in (16,17)
     local io = open(joinpath(@__DIR__, "genStructArrays-v$(v).jl"), "w")
@@ -381,4 +383,6 @@ datatypes = data["datatypes"]
 for (key,value) in pairs(datatypes)
     gen_structarray_rntuple(io, key, value)
 end
+=#
+
 close(io)

src/Components.jl

@@ -24,6 +24,7 @@ Base.iterate(v::Vector3d, i=1) = i > 3 ? nothing : (getproperty(v, propertynames
 Base.length(v::Vector3d) = 3
 θ(v::Vector3d) = atan(√(v.x^2+v.y^2), v.z)
 ϕ(v::Vector3d) = atan(v.y, v.x)
+Base.zero(::Type{Vector3d}) = Vector3d()
 
 #---Vector3f
 Base.show(io::IO, v::Vector3f) = print(io, "($(v.x), $(v.y), $(v.z))")
@@ -40,13 +41,15 @@ Base.iterate(v::Vector3f, i=1) = i > 3 ? nothing : (getproperty(v, propertynames
 Base.length(v::Vector3f) = 3
 θ(v::Vector3f) = atan(√(v.x^2+v.y^2), v.z)
 ϕ(v::Vector3f) = atan(v.y, v.x)
+Base.zero(::Type{Vector3f}) = Vector3f()
 
 #---Vector2i
 Base.show(io::IO, v::Vector2i) = print(io, "($(v.a), $(v.b))")
 Base.:+(v1::Vector2i, v2::Vector2i) = Vector3d(v1.a + v2.a, v1.b + v2.b)
 Base.:-(v1::Vector2i, v2::Vector2i) = Vector3d(v1.a - v2.a, v1.b - v2.b)
 Base.:*(v::Vector2i, a::Number) = Vector3d(a*v.a, a*v.b)
 Base.:*(a::Number, v::Vector2i) = v * a
+Base.zero(::Type{Vector2i}) = Vector2i()
 
 #---Vector4f
 Base.show(io::IO, v::Vector4f) = print(io, "($(v.x), $(v.y), $(v.z), $(v.t))")
@@ -60,6 +63,7 @@ function Base.isapprox(v1::Vector4f, v2::Vector4f; atol::Real=0, rtol::Real=Base
     isapprox(v1.z, v2.z; atol=atol, rtol=rtol, nans=nans) &&
     isapprox(v1.t, v2.t; atol=atol, rtol=rtol, nans=nans)
 end
+Base.zero(::Type{Vector4f}) = Vector4f()
 
 #--------------------------------------------------------------------------------------------------
 #---ObjectID{ED}-----------------------------------------------------------------------------------
@@ -145,6 +149,7 @@ Base.getindex(r::Relation{ED,TD,N}, i) where {ED,TD,N} = 0 < i <= (r.last - r.fi
 Base.size(r::Relation) = (r.last-r.first,)
 Base.length(r::Relation) = r.last-r.first
 Base.eltype(::Type{Relation{ED,TD,N}}) where {ED,TD,N} = TD
+Base.zero(::Type{Relation{ED,TD,N}}) where {ED,TD,N} = Relation{ED,TD,N}(0,0,0)
 function relations(::Type{ED}) where ED
     (ft for ft in fieldtypes(ED) if ft <: Relation)
 end

src/RootIO.jl

@@ -174,8 +174,6 @@ module RootIO
         getproperty(evt, bname)
     end
 
-    include(joinpath(@__DIR__,"../podio/genStructArrays-v16.jl"))
-
     #---Generic StructArray constructor (fall-back)------------------------------------------------
     function StructArray{T,bname}(evt::UnROOT.LazyEvent, collid = UInt32(0), len = -1) where {T,bname} 
         fnames = fieldnames(T)
@@ -191,8 +189,14 @@ module RootIO
                 StructArray{ft,Symbol(bname,:_,fn)}(evt, collid, len)
             elseif ft <: ObjectID              # index of another one....
                 n_rels += 1
-                StructArray{ft, Symbol(bname, "#$(n_rels-1)")}(evt, collid, len)
+                if isnewpodio()
+                    na = replace("$(fn)", "_idx" => "", "mcparticle" => "MCParticle")     # remove the added suffix
+                    StructArray{ft, Symbol("_", bname, "_$(na)")}(evt, collid, len)
+                else
+                    StructArray{ft, Symbol(bname, "#$(n_rels-1)")}(evt, collid, len)
+                end
             else
+                fn == :subdetectorHitNumbers && (fn = :subDetectorHitNumbers)   # adhoc fixes
                 StructArray{ft,Symbol(bname,:_,fn)}(evt, collid, len)
             end
         end)
@@ -289,11 +293,11 @@ module RootIO
         end
         collid = Base.get(reader.collectionIDs, bname, UInt32(0)) # The CollectionID has beeen assigned when opening the file
         sbranch = Symbol(bname)
-        #if reader.isRNTuple
-        #    sa = StructArray{btype}(evt, sbranch, collid)
-        #else
+        if reader.isRNTuple
+            sa = StructArray{btype}(evt, sbranch, collid)
+        else
             sa = StructArray{btype,sbranch}(evt, collid)
-        #end
+        end
         if register
             assignEDStore(sa, collid)
             if !isempty(layout[3])  # check if there are relations in this branch
@@ -320,26 +324,75 @@ module RootIO
         _get(reader, evt, bname, btype, register)
     end
 
-    struct getCollection{ED <: EDM4hep.POD, B}; end
-    function getCollection{ED}(reader::Reader, evt::UnROOT.LazyEvent, bname::String) where ED
-        collid = Base.get(reader.collectionIDs, bname, UInt32(0))
-        StructArray{ED, Symbol(bname)}(evt, collid, -1)
-    end
-    @inline function getCollection{ED,bname}(evt::UnROOT.LazyEvent, collid::UInt32) where {ED,bname}
-        StructArray{ED, bname}(evt, collid)
+    selectednames(btype::Type, selection::Union{AbstractString, Symbol, Regex}) = selectednames(btype, [selection])
+    function selectednames(btype::Type, selection)
+        field_names = fieldnames(btype)
+        _m(r::Regex) = Base.Fix1(occursin, r) ∘ string
+        filtered_names = mapreduce(∪, selection) do b
+            if b isa Regex
+                filter(_m(b), field_names)
+            elseif b isa String
+                Symbol(b) in field_names ? [Symbol(b)] : []
+            elseif b isa Symbol
+                b in field_names ? [b] : []
+            else
+                error("branch selection must be String, Symbol or Regex")
+            end
+        end
+        filtered_names
     end
 
+    """
+    create_getter(reader::Reader, bname::String; selection=nothing)
 
-    function create_getter(reader::Reader, bname::String)
+    This function creates a getter function for a given branch name. The getter function is a function that takes an event and
+    returns a `StructArray` with the data of the branch. The getter function is created as a function with the name `get_<branchname>`.
+    The optional parameter `selection` is a list of field names to be selected from the branch. If `selection` is not provided, all fields are selected.
+    The user can use a list of strings, symbols or regular expressions to select the fields.
+    """
+    function create_getter(reader::Reader, bname::String; selection=nothing)
         btype = reader.btypes[bname]
         collid = Base.get(reader.collectionIDs, bname, UInt32(0))
+        snames = isnothing(selection) ? fieldnames(btype) : selectednames(btype, selection) 
         fname = "get_" * replace(bname, "#" => "_", " " => "_")
-        code = """
-        function $fname(evt::UnROOT.LazyEvent)
-            StructArray{$btype, Symbol(\"$bname\")}(evt, $collid)
+        #---Create the function as a string and evaluate it
+        code = "function $fname(evt::UnROOT.LazyEvent)\n"
+        if btype == ObjectID
+            code *= "    return StructArray{ObjectID}((getproperty(evt, Symbol(\"$bname\",:_index)),getproperty(evt, Symbol(\"$bname\",:_collectionID))))\n"
+        else
+            first = true
+            n_rels = 0    # number of one-to-one or one-to-many Relations
+            for (ft, fn) in zip(fieldtypes(btype), fieldnames(btype))
+                #fn in snames || continue
+                if first
+                    fn == :index && continue    # skip the index field
+                    code *= "    firstmem = getproperty(evt, Symbol(\"$(bname)_$fn\"))\n"
+                    code *= "    len = length(firstmem)\n"
+                    code *= "    columns = (StructArray{ObjectID{$(btype)}}((collect(0:len-1),fill($(collid),len))),\n"
+                    code *= "        firstmem,\n"
+                    first = false
+                else
+                    (ft <: Relation || ft <: ObjectID) && (n_rels += 1)
+                    if fn in snames
+                        if isprimitivetype(ft)
+                            code *= "        getproperty(evt, Symbol(\"$(bname)_$fn\")),\n"
+                        elseif ft <: SVector
+                            N = size(ft)[1]
+                            code *= "        StructArray{$ft}(reshape(getproperty(evt, Symbol(\"$(bname)_$fn[$N]\")), $N, len);dims=1),\n"
+                        elseif ft <: ObjectID  # one-to-one relation
+                            code *= "        StructArray{$ft, Symbol(\"$(bname)#$(n_rels-1)\")}(evt, $(collid), len),\n"
+                        else
+                            code *= "        StructArray{$ft, Symbol(\"$(bname)_$fn\")}(evt, $(collid), len),\n"
+                        end
+                    else
+                        code *= "        zeros($ft, len),\n"
+                    end
+                end
+            end
+            code *= "    )\n"
+            code *= "    return StructArray{$btype}(columns)\n"
         end
-        """
+        code *= "end\n"
         Meta.parse(code) |> eval
     end
-end
-
+end # module RootIO