Commit e7fdff55 authored by David Blyth's avatar David Blyth

Added crude trackres tool for plots of track resolution vs eta and p_T

parent fd9d4b2f
package main
import (
"flag"
"fmt"
"log"
"math"
"os"
"github.com/decibelcooper/proio/go-proio"
"github.com/decibelcooper/proio/go-proio/model/eic"
"go-hep.org/x/hep/hbook"
"gonum.org/v1/plot"
"gonum.org/v1/plot/palette/brewer"
"gonum.org/v1/plot/plotter"
"gonum.org/v1/plot/vg"
"github.com/decibelcooper/eicplot"
)
var (
pTMin = flag.Float64("minpt", 0.5, "minimum transverse momentum")
pTMax = flag.Float64("maxpt", 30, "maximum transverse momentum")
etaLimit = flag.Float64("etalimit", 4, "maximum absolute value of eta")
nBinsPT = flag.Int("nbinspt", 10, "number of bins in transverse momentum")
nBinsEta = flag.Int("nbinseta", 10, "number of bins in eta")
title = flag.String("title", "", "plot title")
output = flag.String("output", "out.png", "output file")
)
func printUsage() {
fmt.Fprintf(os.Stderr, `Usage: `+os.Args[0]+` [options] <proio-input-file>
options:
`,
)
flag.PrintDefaults()
}
func main() {
flag.Usage = printUsage
flag.Parse()
if flag.NArg() != 1 {
printUsage()
log.Fatal("Invalid arguments")
}
p, _ := plot.New()
p.Title.Text = *title
p.X.Label.Text = "eta"
p.Y.Label.Text = "p_T"
p.X.Tick.Marker = eicplot.PreciseTicks{NSuggestedTicks: 5}
p.Y.Tick.Marker = eicplot.PreciseTicks{NSuggestedTicks: 5}
resGrid := NewResGrid(*nBinsEta, -*etaLimit, *etaLimit, *nBinsPT, *pTMin, *pTMax)
filename := flag.Arg(0)
reader, err := proio.Open(filename)
if err != nil {
log.Fatal(err)
}
for event := range reader.ScanEvents() {
ids := event.TaggedEntries("Reconstructed")
for _, id := range ids {
track, ok := event.GetEntry(id).(*eic.Track)
if !ok {
continue
}
partCandID := make(map[uint64]uint64)
for _, obsID := range track.Observation {
eDep, ok := event.GetEntry(obsID).(*eic.EnergyDep)
if !ok {
continue
}
for _, sourceID := range eDep.Source {
simHit, ok := event.GetEntry(sourceID).(*eic.SimHit)
if !ok {
continue
}
partCandID[simHit.Particle]++
}
}
partID := uint64(0)
hitCount := uint64(0)
for id, count := range partCandID {
if count > hitCount {
partID = id
hitCount = count
}
}
part, ok := event.GetEntry(partID).(*eic.Particle)
if !ok {
continue
}
if len(track.Segment) == 0 {
continue
}
pMag := math.Sqrt(math.Pow(part.P.X, 2) + math.Pow(part.P.Y, 2) + math.Pow(part.P.Z, 2))
eta := math.Atanh(part.P.Z / pMag)
pT := math.Sqrt(math.Pow(part.P.X, 2) + math.Pow(part.P.Y, 2))
chargeMag := math.Abs(float64(part.Charge))
poqMag := pMag / chargeMag
trackPoqMag := math.Sqrt(math.Pow(track.Segment[0].Poq.X, 2) +
math.Pow(track.Segment[0].Poq.Y, 2) +
math.Pow(track.Segment[0].Poq.Z, 2))
resGrid.Fill(eta, pT, trackPoqMag/poqMag)
}
}
reader.Close()
pal, _ := brewer.GetPalette(brewer.TypeAny, "RdYlBu", 11)
heatMap := plotter.NewHeatMap(resGrid, pal)
heatMap.Min = 0
heatMap.Max = 0.05
p.Add(heatMap)
p.Save(6*vg.Inch, 4*vg.Inch, *output)
}
type ResGrid struct {
hCount, hV, hV2 *hbook.H2D
nBinsX, nBinsY int
xLow, xHigh float64
yLow, yHigh float64
}
func NewResGrid(nBinsX int, xLow, xHigh float64, nBinsY int, yLow, yHigh float64) *ResGrid {
return &ResGrid{
hbook.NewH2D(nBinsX, xLow, xHigh, nBinsY, yLow, yHigh),
hbook.NewH2D(nBinsX, xLow, xHigh, nBinsY, yLow, yHigh),
hbook.NewH2D(nBinsX, xLow, xHigh, nBinsY, yLow, yHigh),
nBinsX, nBinsY,
xLow, xHigh,
yLow, yHigh,
}
}
func (g *ResGrid) Fill(x, y, z float64) {
g.hCount.Fill(x, y, 1)
g.hV.Fill(x, y, z)
g.hV2.Fill(x, y, z*z)
}
func (g *ResGrid) Dims() (int, int) {
return g.nBinsX, g.nBinsY
}
func (g *ResGrid) Z(i, j int) float64 {
n := g.hCount.GridXYZ().Z(i, j)
if n < 3 {
return 1
}
mean := g.hV.GridXYZ().Z(i, j) / n
mean2 := g.hV2.GridXYZ().Z(i, j) / n
stddev := math.Sqrt(mean2 - mean*mean)
return stddev
}
func (g *ResGrid) X(i int) float64 {
return g.hCount.GridXYZ().X(i)
}
func (g *ResGrid) Y(j int) float64 {
return g.hCount.GridXYZ().Y(j)
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment