main.go 5.49 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
	"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"
)

func printUsage() {
23
	fmt.Fprintf(os.Stderr, `Usage: `+os.Args[0]+` [options] <proio-input-files>...
24 25 26 27 28 29 30 31

options:
`,
	)
	flag.PrintDefaults()
}

func main() {
32 33 34 35 36 37 38 39 40 41 42 43
	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")
44 45
	flag.Usage = printUsage
	flag.Parse()
46
	if flag.NArg() < 1 {
47 48 49 50
		printUsage()
		log.Fatal("Invalid arguments")
	}

51 52 53 54 55
	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}
56

57 58 59 60 61
	nSubs := 1
	nSubs = intMax(nSubs, len(pTMin.Array))
	nSubs = intMax(nSubs, len(pTMax.Array))
	nSubs = intMax(nSubs, len(fracCut.Array))

62
	for i, filename := range flag.Args() {
63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78
		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}
			}
79

80 81 82 83 84 85 86 87 88 89
			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...)
90
		}
91 92 93 94 95 96 97 98
	}

	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)
99

100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116
	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)
117 118 119 120
				if !ok {
					continue
				}

121 122
				for _, sourceID := range eDep.Source {
					simHit, ok := event.GetEntry(sourceID).(*eic.SimHit)
123 124 125 126
					if !ok {
						continue
					}

127
					partCandID[simHit.Particle]++
128
				}
129
			}
130

131 132 133 134 135 136
			partID := uint64(0)
			hitCount := uint64(0)
			for id, count := range partCandID {
				if count > hitCount {
					partID = id
					hitCount = count
137
				}
138
			}
139

140 141 142 143
			part, ok := event.GetEntry(partID).(*eic.Particle)
			if !ok {
				continue
			}
144

145 146 147
			if len(track.Segment) == 0 {
				continue
			}
148

149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164
			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
165 166
			}

167 168
			etaHist.Fill(eta, 1)
		}
169

170 171 172 173 174 175
		ids = event.TaggedEntries("GenStable")
		for _, id := range ids {
			part, ok := event.GetEntry(id).(*eic.Particle)
			if !ok {
				continue
			}
176

177 178 179
			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))
180

181 182 183
			// cuts
			if pT < pTMin || pT > pTMax {
				continue
184 185
			}

186
			trueEtaHist.Fill(eta, 1)
187 188
		}

189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210
		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
211
		}
212 213 214 215
	}
	errPoints := plotutil.ErrorPoints{points, xErrors, yErrors}
	xerr, _ := plotter.NewXErrorBars(errPoints)
	yerr, _ := plotter.NewYErrorBars(errPoints)
216

217 218
	return []plot.Plotter{xerr, yerr}
}
219

220 221 222
func intMin(a, b int) int {
	if a < b {
		return a
223
	}
224 225
	return b
}
226

227 228 229 230 231
func intMax(a, b int) int {
	if a > b {
		return a
	}
	return b
232
}