simple college analysis

event.go

// automatically generated!

package main

import (
	"fmt"
	"math"

	"github.com/go-hep/croot"
	"github.com/go-hep/fmom"
)

type DataReader struct {
	Run      int32
	Event    int32
	Period   int32
	NbElecs  int32
	NbMuons  int32
	L1Charge float32
	L1Type   int32
	L1Pt     float32
	L1Eta    float32
	L1Phi    float32
	L1E      float32
	L2Charge float32
	L2Type   int32
	L2Pt     float32
	L2Eta    float32
	L2Phi    float32
	L2E      float32
	L1L2P    float64

	// branches
	b_Run      croot.Branch
	b_Event    croot.Branch
	b_Period   croot.Branch
	b_NbElecs  croot.Branch
	b_NbMuons  croot.Branch
	b_L1Charge croot.Branch
	b_L1Type   croot.Branch
	b_L1Pt     croot.Branch
	b_L1Eta    croot.Branch
	b_L1Phi    croot.Branch
	b_L1E      croot.Branch
	b_L2Charge croot.Branch
	b_L2Type   croot.Branch
	b_L2Pt     croot.Branch
	b_L2Eta    croot.Branch
	b_L2Phi    croot.Branch
	b_L2E      croot.Branch
	b_L1L2P    croot.Branch

	Tree croot.Tree
}

func NewDataReader(tree croot.Tree) (*DataReader, error) {
	dr := &DataReader{}
	err := dr.Init(tree)
	if err != nil {
		return nil, err
	}
	return dr, nil
}

func (dr *DataReader) Init(tree croot.Tree) error {
	var err error
	var o int32
	dr.Tree = tree

	o = dr.Tree.SetBranchAddress("Run", &dr.Run)
	if o < 0 {
		return fmt.Errorf("invalid branch: [Run] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("Event", &dr.Event)
	if o < 0 {
		return fmt.Errorf("invalid branch: [Event] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("Period", &dr.Period)
	if o < 0 {
		return fmt.Errorf("invalid branch: [Period] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("NbElecs", &dr.NbElecs)
	if o < 0 {
		return fmt.Errorf("invalid branch: [NbElecs] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("NbMuons", &dr.NbMuons)
	if o < 0 {
		return fmt.Errorf("invalid branch: [NbMuons] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L1Charge", &dr.L1Charge)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L1Charge] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L1Type", &dr.L1Type)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L1Type] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L1Pt", &dr.L1Pt)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L1Pt] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L1Eta", &dr.L1Eta)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L1Eta] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L1Phi", &dr.L1Phi)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L1Phi] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L1E", &dr.L1E)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L1E] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L2Charge", &dr.L2Charge)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L2Charge] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L2Type", &dr.L2Type)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L2Type] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L2Pt", &dr.L2Pt)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L2Pt] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L2Eta", &dr.L2Eta)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L2Eta] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L2Phi", &dr.L2Phi)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L2Phi] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L2E", &dr.L2E)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L2E] (got %d)", o)
	}

	o = dr.Tree.SetBranchAddress("L1L2P", &dr.L1L2P)
	if o < 0 {
		return fmt.Errorf("invalid branch: [L1L2P] (got %d)", o)
	}

	return err
}

func (dr *DataReader) GetEntry(entry int64) int {
	if dr.Tree == nil {
		return 0
	}
	return dr.Tree.GetEntry(entry, 1)
}

func (dr *DataReader) L1() Particle {
	pt1 := math.Abs(float64(dr.L1Pt))
	phi1 := float64(dr.L1Phi)
	eta1 := float64(dr.L1Eta)
	ene1 := float64(dr.L1E)
	p1 := fmom.NewPxPyPzE(pt1*math.Cos(phi1), pt1*math.Sin(phi1), pt1*math.Sinh(eta1), ene1)
	return Particle{
		PxPyPzE: &p1,
		Charge:  float64(dr.L1Charge),
		PDG:     int(dr.L1Type),
	}
}

func (dr *DataReader) L2() Particle {
	pt2 := math.Abs(float64(dr.L2Pt))
	phi2 := float64(dr.L2Phi)
	eta2 := float64(dr.L2Eta)
	ene2 := float64(dr.L2E)
	p2 := fmom.NewPxPyPzE(pt2*math.Cos(phi2), pt2*math.Sin(phi2), pt2*math.Sinh(eta2), ene2)
	return Particle{
		PxPyPzE: &p2,
		Charge:  float64(dr.L2Charge),
		PDG:     int(dr.L2Type),
	}
}

func init() {
	// register all generated types with CRoot

}

main.go

package main

import (
	"fmt"
	"image/color"

	"github.com/go-hep/croot"
	"github.com/go-hep/fmom"
	"github.com/go-hep/hbook"
	"github.com/go-hep/hplot"
	"github.com/gonum/plot/vg"
)

type Particle struct {
	*fmom.PxPyPzE
	Charge float64
	PDG    int
}

func abs(v int) int {
	if v < 0 {
		return -v
	}
	return v
}

func main() {
	f, err := croot.OpenFile("FTT-College-periodA.root", "read", "", 1, 0)
	if err != nil {
		panic(err)
	}
	defer f.Close("")

	t := f.Get("FourTopTree").(croot.Tree)
	dr, err := NewDataReader(t)
	if err != nil {
		panic(err)
	}

	nentries := dr.Tree.GetEntries()
	fmt.Printf("entries=%v\n", nentries)

	hee := hbook.NewH1D(100, 0, 180)
	hmm := hbook.NewH1D(100, 0, 180)
	hll := hbook.NewH1D(100, 0, 180)
	hem := hbook.NewH1D(100, 0, 600)

	for i := int64(0); i < nentries; i++ {
		dr.GetEntry(i)

		p1 := dr.L1()
		p2 := dr.L2()
		pll := fmom.Add(p1.PxPyPzE, p2.PxPyPzE)

		if p1.Charge != -p2.Charge {
			continue
		}

		hll.Fill(pll.M(), 1)

		pdg1 := abs(p1.PDG)
		pdg2 := abs(p2.PDG)

		switch {
		case pdg1 == 11 && pdg2 == 11:
			hee.Fill(pll.M(), 1)
		case pdg1 == 13 && pdg2 == 13:
			hmm.Fill(pll.M(), 1)
		case pdg1 != pdg2:
			hem.Fill(pll.M(), 1)
			m := pll.M()
			if m < 0 || m > 600 {
				fmt.Printf(">>> %v GeV/c^2\n", m)
			}
		}
	}

	fmt.Printf("hll: entries=%v mean=%v rms=%v\n", hll.Entries(), hll.Mean(), hll.RMS())
	fmt.Printf("hee: entries=%v mean=%v rms=%v\n", hee.Entries(), hee.Mean(), hee.RMS())
	fmt.Printf("hmm: entries=%v mean=%v rms=%v\n", hmm.Entries(), hmm.Mean(), hmm.RMS())
	fmt.Printf("hem: entries=%v mean=%v rms=%v (%v, %v)\n", hem.Entries(), hem.Mean(), hem.RMS(),
		hem.Annotation()["underflow"], hem.Annotation()["overflow"],
	)
	fmt.Printf("hem: %v\n", hem.Annotation())

	p, err := hplot.New()
	if err != nil {
		panic(err)
	}
	p.Title.Text = "l+l- invariant mass"
	p.X.Label.Text = "Invariant Mass (GeV/c^2)"
	p.X.Min = 0
	p.X.Max = 180

	phll, err := hplot.NewH1D(hll)
	if err != nil {
		panic(err)
	}
	p.Add(phll)

	if true {
		hplee, err := hplot.NewH1D(hee)
		if err != nil {
			panic(err)
		}
		hplee.Color = color.RGBA{255, 0, 0, 255}

		hplmm, err := hplot.NewH1D(hmm)
		if err != nil {
			panic(err)
		}
		hplmm.Color = color.RGBA{0, 0, 255, 255}

		p.Add(hplee, hplmm)
	}

	err = p.Save(20*vg.Centimeter, 10*vg.Centimeter, "plot.png")
	if err != nil {
		panic(err)
	}

	{
		p, err := hplot.New()
		if err != nil {
			panic(err)
		}
		p.Title.Text = "e-mu invariant mass"
		p.X.Label.Text = "Invariant mass (GeV/c^2)"

		hplem, err := hplot.NewH1D(hem)
		if err != nil {
			panic(err)
		}
		hplem.Infos.Style = hplot.HInfoSummary

		p.Add(hplem)
		err = p.Save(20*vg.Centimeter, 10*vg.Centimeter, "plot-hem.png")
		if err != nil {
			panic(err)
		}
	}
}

A simple analysis for College (junior high)