main.go 4.57 KB
Newer Older
1 2 3 4 5
package main

import (
	"flag"
	"fmt"
6
	"image/color"
7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27
	"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/plotter"
	"gonum.org/v1/plot/plotutil"
	"gonum.org/v1/plot/vg"

	"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")
28 29
	title    = flag.String("title", "", "plot title")
	prefix   = flag.String("prefix", "out", "output file prefix")
30 31 32
)

func printUsage() {
33
	fmt.Fprintf(os.Stderr, `Usage: `+os.Args[0]+` [options] <proio-input-files>...
34 35 36 37 38 39 40 41 42 43

options:
`,
	)
	flag.PrintDefaults()
}

func main() {
	flag.Usage = printUsage
	flag.Parse()
44
	if flag.NArg() < 1 {
45 46 47 48
		printUsage()
		log.Fatal("Invalid arguments")
	}

49 50 51 52 53
	p, _ := plot.New()
	p.Title.Text = *title
	p.X.Label.Text = "eta"
	p.X.Tick.Marker = eicplot.PreciseTicks{NSuggestedTicks: 5}
	p.Y.Tick.Marker = eicplot.PreciseTicks{NSuggestedTicks: 5}
54

55 56 57 58 59 60 61 62
	for i, filename := range flag.Args() {
		etaHist := hbook.NewH1D(*nBins, -*etaLimit, *etaLimit)
		trueEtaHist := hbook.NewH1D(*nBins, -*etaLimit, *etaLimit)

		reader, err := proio.Open(filename)
		if err != nil {
			log.Fatal(err)
		}
63

64 65 66 67 68
		eventNum := 0
		for event := range reader.ScanEvents() {
			ids := event.TaggedEntries("Reconstructed")
			for _, id := range ids {
				track, ok := event.GetEntry(id).(*eic.Track)
69 70 71 72
				if !ok {
					continue
				}

73 74 75
				partCandID := make(map[uint64]uint64)
				for _, obsID := range track.Observation {
					eDep, ok := event.GetEntry(obsID).(*eic.EnergyDep)
76 77 78 79
					if !ok {
						continue
					}

80 81 82 83 84 85 86 87
					for _, sourceID := range eDep.Source {
						simHit, ok := event.GetEntry(sourceID).(*eic.SimHit)
						if !ok {
							continue
						}

						partCandID[simHit.Particle]++
					}
88 89
				}

90 91 92 93 94 95 96
				partID := uint64(0)
				hitCount := uint64(0)
				for id, count := range partCandID {
					if count > hitCount {
						partID = id
						hitCount = count
					}
97 98
				}

99 100 101 102
				part, ok := event.GetEntry(partID).(*eic.Particle)
				if !ok {
					continue
				}
103

104 105 106
				if len(track.Segment) == 0 {
					continue
				}
107

108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126
				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
				}

				etaHist.Fill(eta, 1)
127 128
			}

129 130 131 132 133 134
			ids = event.TaggedEntries("GenStable")
			for _, id := range ids {
				part, ok := event.GetEntry(id).(*eic.Particle)
				if !ok {
					continue
				}
135

136 137 138
				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))
139

140 141 142 143
				// cuts
				if pT < *pTMin {
					continue
				}
144

145
				trueEtaHist.Fill(eta, 1)
146 147
			}

148
			eventNum++
149 150
		}

151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179
		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}
180
		}
181 182 183 184 185 186 187
		xerr.LineStyle.Color = pointColor
		yerr.LineStyle.Color = pointColor

		p.Add(xerr)
		p.Add(yerr)

		reader.Close()
188 189
	}

190 191
	p.Save(6*vg.Inch, 4*vg.Inch, *prefix+".pdf")
	p.Save(6*vg.Inch, 4*vg.Inch, *prefix+".png")
192
}