DisplacedMuonGun

import ROOT
import numpy as n

f = ROOT.TFile("flatDxyDistribution.root", "recreate")
t = ROOT.TTree("name_of_tree", "tree title")


# create 1 dimensional float arrays (python's float datatype corresponds to c++ doubles)
# as fill variables
pt = n.zeros(1, dtype=float)
px = n.zeros(1, dtype=float)
py = n.zeros(1, dtype=float)
pz = n.zeros(1, dtype=float)
p = n.zeros(1, dtype=float)
m = n.zeros(1, dtype=float)
vx = n.zeros(1, dtype=float)
vy = n.zeros(1, dtype=float)
vz = n.zeros(1, dtype=float)
phi = n.zeros(1, dtype=float)
eta = n.zeros(1, dtype=float)
dxy = n.zeros(1, dtype=float)
lxy = n.zeros(1, dtype=float)
#check position in st1 
x_st1 = n.zeros(1, dtype=float)
y_st1 = n.zeros(1, dtype=float)
z_st1 = n.zeros(1, dtype=float)
r_st1 = n.zeros(1, dtype=float)
phi_st1 = n.zeros(1, dtype=float)
eta_st1 = n.zeros(1, dtype=float)
h_r = ROOT.TH1F("h_r","h_r",1000, 0, 1000)


# create the branches and assign the fill-variables to them
t.Branch('px', px, 'px/D')
t.Branch('py', py, 'py/D')
t.Branch('pz', pz, 'pz/D')
t.Branch('p', p, 'p/D')
t.Branch('pt', pt, 'pt/D')
t.Branch('vx', vx, 'vx/D')
t.Branch('vy', vy, 'vy/D')
t.Branch('vz', vz, 'vz/D')
t.Branch('dxy', dxy, 'dxy/D')
t.Branch('lxy', lxy, 'lxy/D')
t.Branch('phi', phi, 'phi/D')
t.Branch('eta', eta, 'eta/D')

t.Branch('x_st1', x_st1, 'x_st1/D')
t.Branch('y_st1', y_st1, 'y_st1/D')
t.Branch('z_st1', z_st1, 'z_st1/D')
t.Branch('r_st1', r_st1, 'r_st1/D')
t.Branch('phi_st1', phi_st1, 'phi_st1/D')
t.Branch('eta_st1', eta_st1, 'eta_st1/D')


# create some random numbers, fill them into the fill varibles and call Fill()
for i in xrange(10000):
    pt[0] = ROOT.gRandom.Uniform(2.,50.)
    dxy[0] = ROOT.gRandom.Uniform(0.,50.)
    phi[0] = ROOT.gRandom.Uniform(0,2*n.pi)
    
    ## px and py
    px[0] = pt[0] * n.cos(phi[0])
    py[0] = pt[0] * n.sin(phi[0])

    getlxy = False 
    for i in range(0, 10000):
	vx[0] = ROOT.gRandom.Uniform(-50,50)
	vy[0] = (pt[0]*dxy[0] + vx[0] * py[0])/px[0]
	lxy[0] = n.sqrt(vx[0]*vx[0] + vy[0]*vy[0])
	if lxy[0] < 50:
	    getlxy = True
	    break
    if not(getlxy):
	continue

    eta[0] = ROOT.gRandom.Uniform(-2.5, 2.5)
    pz[0] = pt[0]*n.sinh(eta[0])
    p[0] = n.sqrt(px[0]*px[0] + py[0]*py[0] + pz[0]*pz[0])
    #print "pt ",pt[0] ," pz ",pz[0]," p[0] ",p[0]," pt*cosh(eta) ", pt[0]*n.cosh(eta[0])
	    #eta[0] = 0.5*n.log((p[0]+pz[0])/(p[0]-pz[0]))

    vz[0] = ROOT.gRandom.Uniform(-500, 500)
    #line function (x, y , z) = (vx, vy, vz) + alpha*(px, py, pz)
    z_st1[0] = 600
    alpha = (z_st1[0]-vz[0])/pz[0]
    x_st1[0] = vx[0] + alpha*px[0]
    y_st1[0] = vy[0] + alpha*py[0]
    print "vx ",vx[0]," vy ",vy[0]," vz ",vz[0]," Z ",z_st1[0]," alpha ",alpha, " X ",x_st1[0]," Y ",y_st1[0]
    r_st1[0] = n.sqrt(x_st1[0]*x_st1[0] + y_st1[0]*y_st1[0])
    h_r.Fill(r_st1[0])
    phi_st1[0] = n.arctan(y_st1[0]/x_st1[0])

    t.Fill()

# write the tree into the output file and close the file
f.Write()
f.Close()