trackeff: improved ability to overlay multiple plots

flags.go

+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)
+}

trackeff/main.go

@@ -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
 }