Commit 1e63faaf authored by David Blyth's avatar David Blyth

Added trackpull, a companion to trackres

parent d6b6b962
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/moreland"
"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 := NewPullGrid(*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()
colorMap := moreland.SmoothBlueRed()
colorMap.SetMin(0.99)
colorMap.SetMax(1.01)
heatMap := plotter.NewHeatMap(resGrid, colorMap.Palette(1000))
heatMap.Min = 0.99
heatMap.Max = 1.01
p.Add(heatMap)
p.Save(6*vg.Inch, 4*vg.Inch, *output)
p, _ = plot.New()
colorBar := &plotter.ColorBar{ColorMap: colorMap}
colorBar.Vertical = true
p.Add(colorBar)
p.HideX()
p.Y.Padding = 0
p.Y.Tick.Marker = eicplot.PreciseTicks{NSuggestedTicks: 3}
p.Save(50, 400, "scale.png")
}
type PullGrid struct {
hCount, hV, hV2 *hbook.H2D
nBinsX, nBinsY int
xLow, xHigh float64
yLow, yHigh float64
}
func NewPullGrid(nBinsX int, xLow, xHigh float64, nBinsY int, yLow, yHigh float64) *PullGrid {
return &PullGrid{
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 *PullGrid) 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 *PullGrid) Dims() (int, int) {
return g.nBinsX, g.nBinsY
}
func (g *PullGrid) Z(i, j int) float64 {
n := g.hCount.GridXYZ().Z(i, j)
if n < 3 {
return 0
}
mean := g.hV.GridXYZ().Z(i, j) / n
return mean
}
func (g *PullGrid) X(i int) float64 {
return g.hCount.GridXYZ().X(i)
}
func (g *PullGrid) 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