Commit 12f53238 authored by David Blyth's avatar David Blyth

trackeff: improved ability to overlay multiple plots

parent eca71dfa
package eicplot
import (
"fmt"
"strconv"
)
type FloatArrayFlags struct {
Array []float64
beenSet bool
}
func (f *FloatArrayFlags) Set(valueStr string) error {
value, err := strconv.ParseFloat(valueStr, 64)
if err != nil {
return err
}
if !f.beenSet {
f.beenSet = true
f.Array = nil
}
f.Array = append(f.Array, value)
return nil
}
func (f *FloatArrayFlags) String() string {
return fmt.Sprint(f.Array)
}
......@@ -19,16 +19,6 @@ import (
"github.com/decibelcooper/eicplot"
)
var (
pTMin = flag.Float64("minpt", 0.5, "minimum transverse momentum")
pTMax = flag.Float64("maxpt", 100, "maximum transverse momentum")
fracCut = flag.Float64("frac", 0.01, "maximum fractional magnitude of the difference in momentum between track and true")
etaLimit = flag.Float64("etalimit", 4, "maximum absolute value of eta")
nBins = flag.Int("nbins", 80, "number of bins")
title = flag.String("title", "", "plot title")
prefix = flag.String("prefix", "out", "output file prefix")
)
func printUsage() {
fmt.Fprintf(os.Stderr, `Usage: `+os.Args[0]+` [options] <proio-input-files>...
......@@ -39,6 +29,18 @@ options:
}
func main() {
pTMin := &eicplot.FloatArrayFlags{Array: []float64{0.5}}
pTMax := &eicplot.FloatArrayFlags{Array: []float64{100000}}
fracCut := &eicplot.FloatArrayFlags{Array: []float64{0.01}}
var (
etaLimit = flag.Float64("etalimit", 4, "maximum absolute value of eta")
nBins = flag.Int("nbins", 80, "number of bins")
title = flag.String("title", "", "plot title")
output = flag.String("output", "out.png", "output file")
)
flag.Var(pTMin, "minpt", "minimum transverse momentum")
flag.Var(pTMax, "maxpt", "maximum transverse momentum")
flag.Var(fracCut, "frac", "maximum fractional magnitude of the difference in momentum between track and true")
flag.Usage = printUsage
flag.Parse()
if flag.NArg() < 1 {
......@@ -52,141 +54,179 @@ func main() {
p.X.Tick.Marker = eicplot.PreciseTicks{NSuggestedTicks: 5}
p.Y.Tick.Marker = eicplot.PreciseTicks{NSuggestedTicks: 5}
nSubs := 1
nSubs = intMax(nSubs, len(pTMin.Array))
nSubs = intMax(nSubs, len(pTMax.Array))
nSubs = intMax(nSubs, len(fracCut.Array))
for i, filename := range flag.Args() {
etaHist := hbook.NewH1D(*nBins, -*etaLimit, *etaLimit)
trueEtaHist := hbook.NewH1D(*nBins, -*etaLimit, *etaLimit)
for j := 0; j < nSubs; j++ {
iPTMin := intMin(j, len(pTMin.Array)-1)
iPTMax := intMin(j, len(pTMax.Array)-1)
iFracCut := intMin(j, len(fracCut.Array)-1)
plotters := makeTrackEffPlotters(filename, pTMin.Array[iPTMin], pTMax.Array[iPTMax], fracCut.Array[iFracCut], *etaLimit, *nBins)
pointColor := color.RGBA{A: 255}
switch i + j {
case 1:
pointColor = color.RGBA{G: 255, A: 255}
case 2:
pointColor = color.RGBA{B: 255, A: 255}
case 3:
pointColor = color.RGBA{R: 255, B: 127, G: 127, A: 255}
}
reader, err := proio.Open(filename)
if err != nil {
log.Fatal(err)
for _, p := range plotters {
switch t := p.(type) {
case *plotter.XErrorBars:
t.LineStyle.Color = pointColor
case *plotter.YErrorBars:
t.LineStyle.Color = pointColor
}
}
p.Add(plotters...)
}
}
p.Save(6*vg.Inch, 4*vg.Inch, *output)
}
func makeTrackEffPlotters(filename string, pTMin, pTMax, fracCut, etaLimit float64, nBins int) []plot.Plotter {
etaHist := hbook.NewH1D(nBins, -etaLimit, etaLimit)
trueEtaHist := hbook.NewH1D(nBins, -etaLimit, etaLimit)
eventNum := 0
for event := range reader.ScanEvents() {
ids := event.TaggedEntries("Reconstructed")
for _, id := range ids {
track, ok := event.GetEntry(id).(*eic.Track)
reader, err := proio.Open(filename)
if err != nil {
log.Fatal(err)
}
eventNum := 0
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
}
partCandID := make(map[uint64]uint64)
for _, obsID := range track.Observation {
eDep, ok := event.GetEntry(obsID).(*eic.EnergyDep)
for _, sourceID := range eDep.Source {
simHit, ok := event.GetEntry(sourceID).(*eic.SimHit)
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
}
partCandID[simHit.Particle]++
}
}
part, ok := event.GetEntry(partID).(*eic.Particle)
if !ok {
continue
partID := uint64(0)
hitCount := uint64(0)
for id, count := range partCandID {
if count > hitCount {
partID = id
hitCount = count
}
}
if len(track.Segment) == 0 {
continue
}
part, ok := event.GetEntry(partID).(*eic.Particle)
if !ok {
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
diffMag := math.Sqrt(math.Pow(track.Segment[0].Poq.X-part.P.X/chargeMag, 2) +
math.Pow(track.Segment[0].Poq.Y-part.P.Y/chargeMag, 2) +
math.Pow(track.Segment[0].Poq.Z-part.P.Z/chargeMag, 2))
fracDiff := diffMag / poqMag
// cuts
if pT < *pTMin {
continue
}
if fracDiff > *fracCut {
continue
}
if len(track.Segment) == 0 {
continue
}
etaHist.Fill(eta, 1)
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
diffMag := math.Sqrt(math.Pow(track.Segment[0].Poq.X-part.P.X/chargeMag, 2) +
math.Pow(track.Segment[0].Poq.Y-part.P.Y/chargeMag, 2) +
math.Pow(track.Segment[0].Poq.Z-part.P.Z/chargeMag, 2))
fracDiff := diffMag / poqMag
// cuts
if pT < pTMin || pT > pTMax {
continue
}
if fracDiff > fracCut {
continue
}
ids = event.TaggedEntries("GenStable")
for _, id := range ids {
part, ok := event.GetEntry(id).(*eic.Particle)
if !ok {
continue
}
etaHist.Fill(eta, 1)
}
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))
ids = event.TaggedEntries("GenStable")
for _, id := range ids {
part, ok := event.GetEntry(id).(*eic.Particle)
if !ok {
continue
}
// cuts
if pT < *pTMin {
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))
trueEtaHist.Fill(eta, 1)
// cuts
if pT < pTMin || pT > pTMax {
continue
}
eventNum++
trueEtaHist.Fill(eta, 1)
}
points := make(plotter.XYs, *nBins)
xErrors := make(plotter.XErrors, *nBins)
yErrors := make(plotter.YErrors, *nBins)
binHalfWidth := *etaLimit / float64(*nBins)
binSigma := binHalfWidth / math.Sqrt(3.)
for i := range points {
trueX, trueY := trueEtaHist.XY(i)
points[i].X = trueX + binHalfWidth
xErrors[i].Low = binSigma
xErrors[i].High = binSigma
_, trackY := etaHist.XY(i)
if trueY > 0 {
points[i].Y = trackY / trueY
yErrors[i].Low = math.Sqrt((1 - trackY/trueY) * trackY / math.Pow(trueY, 2))
yErrors[i].High = yErrors[i].Low
}
}
errPoints := plotutil.ErrorPoints{points, xErrors, yErrors}
xerr, _ := plotter.NewXErrorBars(errPoints)
yerr, _ := plotter.NewYErrorBars(errPoints)
pointColor := color.RGBA{A: 255}
switch i {
case 1:
pointColor = color.RGBA{G: 255, A: 255}
case 2:
pointColor = color.RGBA{B: 255, A: 255}
eventNum++
}
reader.Close()
points := make(plotter.XYs, nBins)
xErrors := make(plotter.XErrors, nBins)
yErrors := make(plotter.YErrors, nBins)
binHalfWidth := etaLimit / float64(nBins)
binSigma := binHalfWidth / math.Sqrt(3.)
for i := range points {
trueX, trueY := trueEtaHist.XY(i)
points[i].X = trueX + binHalfWidth
xErrors[i].Low = binSigma
xErrors[i].High = binSigma
_, trackY := etaHist.XY(i)
if trueY > 0 {
points[i].Y = trackY / trueY
yErrors[i].Low = math.Sqrt((1 - trackY/trueY) * trackY / math.Pow(trueY, 2))
yErrors[i].High = yErrors[i].Low
}
xerr.LineStyle.Color = pointColor
yerr.LineStyle.Color = pointColor
}
errPoints := plotutil.ErrorPoints{points, xErrors, yErrors}
xerr, _ := plotter.NewXErrorBars(errPoints)
yerr, _ := plotter.NewYErrorBars(errPoints)
p.Add(xerr)
p.Add(yerr)
return []plot.Plotter{xerr, yerr}
}
reader.Close()
func intMin(a, b int) int {
if a < b {
return a
}
return b
}
p.Save(6*vg.Inch, 4*vg.Inch, *prefix+".pdf")
p.Save(6*vg.Inch, 4*vg.Inch, *prefix+".png")
func intMax(a, b int) int {
if a > b {
return a
}
return b
}
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