Fix plotting for ClusterSequence and update serialisation

README.md

@@ -86,12 +86,14 @@ The example also shows how to use `JetReconstruction.HepMC3` to read HepMC3 ASCI
 
 ### Plotting
 
-**TO BE FIXED**
-
 ![illustration](img/illustration.jpeg)
 
 To visualise the clustered jets as a 3d bar plot (see illustration above) we now use `Makie.jl`. See the `jetsplot` function and its documentation for more. 
 
+### Serialisation
+
+The package also provides methods such as `loadjets`, `loadjets!`, and `savejets` that one can use to save and load objects on/from disk easily in a very flexible format. See documentation for more.
+
 ## Reference
 
 Although it has been developed further since the CHEP2023 conference, the CHEP conference proceedings, [arXiv:2309.17309](https://arxiv.org/abs/2309.17309), should be cited if you use this package:

src/JetVis.jl

@@ -1,10 +1,61 @@
 ## Jet visualisation
 # not a submodule
 
+function get_all_ancestors(idx, cs::ClusterSequence)
+	if cs.history[idx].parent1 == NonexistentParent
+		return [cs.history[idx].jetp_index]
+	#elseif cs.history[idx].parent2 == BeamJet
+	#	return 
+	else
+		branch1 = get_all_ancestors(cs.history[idx].parent1, cs)
+		cs.history[idx].parent2 == BeamJet && return branch1
+		branch2 = get_all_ancestors(cs.history[idx].parent2, cs)
+		return [branch1; branch2]
+	end
+end
+
+"""
+    jetsplot(objects, cs::ClusterSequence; barsize_phi=0.1, barsize_eta=0.1, colormap=:glasbey_hv_n256, Module=Main)
+
+Plots a 3d bar chart that represents jets. Takes `objects`, an array of objects to display (should be the same array you have passed to `jet_reconstruct` to get the `cs::ClusterSequence`), and the `cs::ClusterSequence` itself as arguments.
+
+Optional arguments:
+`barsize_phi::Real` — width of a bar along the ϕ axis;
+`barsize_eta::Real` — width of a bar along the η axis;
+`colormap::Symbol` — Makie colour map;
+`Module` — the module where you have your Makie (see below);
+```
+# example
+using CairoMakie # use any other Makie that you have here
+jetsplot([object1, object2, object3], cluster_sequence_I_got_from_jet_reconstruct; Module=CairoMakie)
+```
+
+This function needs `Makie.jl` to work. You should install and import/use a specific backend yourself. `jetsplot` works with `CairoMakie`, `WGLMakie`, `GLMakie`, etc. Additionally, you can specify the module where you have your `Makie` explicitly:
+```
+import CairoMakie
+jetsplot(my_objects, cs, Module=CairoMakie)
+
+import GLMakie
+jetsplot(my_objects, cs, Module=GLMakie)
+
+using WGLMakie
+jetsplot(my_objects, cs, Module=Main) #default
+```
+"""
+function jetsplot(objects, cs::ClusterSequence; barsize_phi=0.1, barsize_eta=0.1, colormap=:glasbey_hv_n256, Module=Main)
+	idx_arrays = Vector{Int}[]
+    for elt in cs.history
+        elt.parent2 == BeamJet || continue
+		push!(idx_arrays, get_all_ancestors(elt.parent1, cs))
+    end
+
+	jetsplot(objects, idx_arrays; barsize_phi, barsize_eta, colormap, Module)
+end
+
 """
 `jetsplot(objects, idx_arrays; barsize_phi=0.1, barsize_eta=0.1, colormap=:glasbey_hv_n256, Module=Main)`
 
-Plots a 3d bar chart that represents jets. Takes an `objects` array of objects to display and `idx_arrays`, an array of arrays with indeces, where `idx_arrays[i]` gives indeces of `objects` that form the jet number `i`.
+Plots a 3d bar chart that represents jets. Takes an `objects` array of objects to display and `idx_arrays`, an array of arrays with indeces, where `idx_arrays[i]` gives indeces of `objects` that form the jet number `i`. This function's signature might not be the most practical for the current version of the JetReconstruction.jl package, as it has been written during the early stage of development. There is now an overload of it that takes a `ClusterSequence` object as its argument.
 
 Optional arguments:
 `barsize_phi::Real` — width of a bar along the ϕ axis;
@@ -32,15 +83,15 @@ jetsplot(my_objects, my_colour_arrays, Module=Main) #default
 ```
 """
 function jetsplot(objects, idx_arrays; barsize_phi=0.1, barsize_eta=0.1, colormap=:glasbey_hv_n256, Module=Main)
-	cs = fill(0, length(objects))
+	cs = fill(0, length(objects)) # colours
 	for i in 1:length(idx_arrays), j in idx_arrays[i]
 		cs[j] = i
 	end
 
-	pts = pt.(objects)
+	pts = sqrt.(pt2.(objects))
 
 	Module.meshscatter(
-		Module.Point3f.(phi.(objects), eta.(objects), 0pts);
+		Module.Point3f.(phi.(objects), rapidity.(objects), 0pts);
 	  	color = cs,
 		markersize = Module.Vec3f.(barsize_phi, barsize_eta, pts),
 		colormap = colormap,

src/Serialize.jl

@@ -1,19 +1,18 @@
 """
 `savejets(filename, jets; format="px py pz E")`
 
-Saves the given `jets` into a file with the given `filename`. Each line contains information about a single jet and is formatted according to the `format` string which defaults to `"E px py pz"` but can also contain other values in any order: `"pt"` or `"kt"` for transverse momentum, `"phi"` for azimuth, `"eta"` for pseudorapidity, `"m"` for mass. It is strongly NOT recommended to put something other than values and (possibly custom) separators in the `format` string.
+Saves the given `jets` into a file with the given `filename`. Each line contains information about a single jet and is formatted according to the `format` string which defaults to `"px py pz E"` but can also contain other values in any order: `"pt2"` for pt^2, `"phi"` for azimuth, `"rapidity"` for rapidity. It is strongly NOT recommended to put something other than values and (possibly custom) separators in the `format` string.
 """
 function savejets(filename, jets; format="px py pz E")
     symbols = Dict(
         "E" => JetReconstruction.energy,
+        "energy" => JetReconstruction.energy,
         "px" => JetReconstruction.px,
-        "pt" => JetReconstruction.pt,
-        "kt" => JetReconstruction.pt,
         "py" => JetReconstruction.py,
         "pz" => JetReconstruction.pz,
+        "pt2" => JetReconstruction.pt2,
         "phi" => JetReconstruction.phi,
-        "m" => JetReconstruction.mass,
-        "eta" => JetReconstruction.eta
+        "rapidity" => JetReconstruction.rapidity
     )
     for pair in symbols
         if !occursin(pair[1], format)
@@ -22,7 +21,7 @@ function savejets(filename, jets; format="px py pz E")
     end
 
     open(filename, "w") do file
-        write(file, "# this file contains jet data.\n# each line that does not start with '#' contains the information about a jet in the following format:\n# "*"$(format)"*"\n# where E is energy, px is momentum along x, py is momentum along y, pz is momentum along z, pt or kt is transverse momentum, phi is azimuth, eta is pseudorapidity, m is mass\n")
+        write(file, "# this file contains jet data.\n# each line that does not start with '#' contains the information about a jet in the following format:\n# "*"$(format)"*"\n# where E is energy, px is momentum along x, py is momentum along y, pz is momentum along z, pt2 is pt^2, phi is azimuth, and rapidity is rapidity\n")
         for j in jets
             line = format
             for pair in symbols
@@ -36,9 +35,9 @@ function savejets(filename, jets; format="px py pz E")
 end
 
 """
-`loadjets!(filename, jets; splitby=isspace, constructor=(px,py,pz,E)->[px,py,pz,E], dtype=Float64) -> jets`
+    loadjets!(filename, jets; splitby=isspace, constructor=(px,py,pz,E)->LorentzVector(E,px,py,pz), dtype=Float64)
 
-Loads the `jets` from a file. Ignores lines that start with `'#'`. Each line gets processed in the following way: the line is split using `split(line, splitby)` or simply `split(line)` by default. Every value in this line is then converted to the `dtype` (which is `Float64` by default). These values are then used as arguments for the `constructor` function which should produce individual jets. By default, the `constructor` constructs vectors of the form `[px,py,pz,E]`.
+Loads the `jets` from a file. Ignores lines that start with `'#'`. Each line gets processed in the following way: the line is split using `split(line, splitby)` or simply `split(line)` by default. Every value in this line is then converted to the `dtype` (which is `Float64` by default). These values are then used as arguments for the `constructor` function which should produce individual jets. By default, the `constructor` constructs Lorentz vectors.
 
 Everything that was already in `jets` is not affected as we only use `push!` on it.
 ```julia
@@ -48,7 +47,7 @@ loadjets!("myjets1.dat", jets)
 loadjets!("myjets2.dat", jets)
 ```
 """
-function loadjets!(filename, jets; splitby=isspace, constructor=(x,y,z,E)->Float64[x,y,z,E], dtype=Float64)
+function loadjets!(filename, jets; splitby=isspace, constructor=(px,py,pz,E)->LorentzVector(E,px,py,pz), dtype=Float64)
     open(filename, "r") do file
         for line in eachline(file)
             if line[1] != '#'
@@ -64,15 +63,15 @@ function loadjets!(filename, jets; splitby=isspace, constructor=(x,y,z,E)->Float
 end
 
 """
-`loadjets(filename; splitby=isspace, constructor=(px,py,pz,E)->Float64[px,py,pz,E], dtype=Float64) -> jets`
+    loadjets(filename; splitby=isspace, constructor=(px,py,pz,E)->LorentzVector(E,px,py,pz), VT=LorentzVector) -> jets
 
-Loads the `jets` from a file. Ignores lines that start with `'#'`. Each line gets processed in the following way: the line is split using `split(line, splitby)` or simply `split(line)` by default. Every value in this line is then converted to the `dtype` (which is `Float64` by default). These values are then used as arguments for the `constructor` function which should produce individual jets. By default, the `constructor` constructs vectors of the form `[px,py,pz,E]`.
+Loads the `jets` from a file, where each element of `jets` is of type `VT`. Ignores lines that start with `'#'`. Each line gets processed in the following way: the line is split using `split(line, splitby)` or simply `split(line)` by default. These values are then used as arguments for the `constructor` function which should produce individual jets of the `VT` type. By default, the `constructor` constructs Lorentz vectors.
 
 ```julia
 # example
 jets = loadjets("myjets1.dat")
 ```
 """
-function loadjets(filename; splitby=isspace, constructor=(px,py,pz,E)->Float64[px,py,pz,E], dtype=Float64)
-    loadjets!(filename, Vector{Float64}[], splitby=splitby, constructor=constructor)
+function loadjets(filename; splitby=isspace, constructor=(px,py,pz,E)->LorentzVector(E,px,py,pz), VT=LorentzVector)
+    loadjets!(filename, VT[], splitby=splitby, constructor=constructor)
 end