Commit daa5ceb7 authored by janlukas.bosse's avatar janlukas.bosse
Browse files

asdf

parent c5f92544
import sys
import matplotlib.pyplot as plt
class Bin1D:
def __init__(self, xmin, xmax):
self.xmin = xmin
self.xmax = xmax
self.w = 0.
self.w2 = 0.
self.wx = 0.
self.wx2 = 0.
self.n = 0.
def __repr__(self):
return str(self)
def __str__(self):
return "{0:10.6e}\t{1:10.6e}\t{2:10.6e}\t{3:10.6e}\t{4:10.6e}\t{5:10.6e}\t{6}".format(
self.xmin, self.xmax, self.w, self.w2, self.wx, self.wx2, int(self.n))
def Format(self, tag):
return "{0}\t{0}\t{1:10.6e}\t{2:10.6e}\t{3:10.6e}\t{4:10.6e}\t{5}".format(
tag, self.w, self.w2, self.wx, self.wx2, int(self.n))
def Fill(self, x, w):
self.w += w
self.w2 += w * w
self.wx += w * x
self.wx2 += w * w * x
self.n += 1.
def ScaleW(self, scale):
self.w *= scale
self.w2 *= scale * scale
self.wx *= scale
self.wx2 *= scale * scale
class Point2D:
def __init__(self, xmin, xmax):
self.x = (xmax + xmin) / 2
self.xerr = [(xmax - xmin) / 2] * 2
self.y = 0
self.yerr = [0] * 2
def __repr__(self):
return str(self)
def __str__(self):
return "{0:10.6e}\t{1:10.6e}\t{2:10.6e}\t{3:10.6e}\t{4:10.6e}\t{5:10.6e}".format(
self.x, self.xerr[0], self.xerr[1], self.y, self.yerr[0], self.yerr[1])
def ScaleY(self, scale):
self.y *= scale
self.yerr = [err * scale for err in self.yerr]
class Histo1D:
def __init__(self, nbin, xmin, xmax, name="/MC/untitled"):
self.name = name
self.bins = []
self.uflow = Bin1D(-sys.float_info.max, xmin)
width = (xmax - xmin) / nbin
for i in range(nbin):
self.bins.append(Bin1D(xmin + i * width, xmin + (i + 1) * width))
self.oflow = Bin1D(xmax, sys.float_info.max)
self.total = Bin1D(-sys.float_info.max, sys.float_info.max)
self.scale = 1.
def __repr__(self):
return str(self)
def __str__(self):
s = "# BEGIN YODA_HISTO1D {0}\n".format(self.name)
s += "Path={0}\n".format(self.name)
s += "ScaledBy={0}\n".format(self.scale)
s += "Title=\nType=Histo1D\n"
s += "# ID\tID\tsumw\tsumw2\tsumwx\tsumwx2\tnumEntries\n"
s += self.total.Format("Total") + "\n"
s += self.uflow.Format("Underflow") + "\n"
s += self.oflow.Format("Overflow") + "\n"
s += "# xlow\txhigh\tsumw\tsumw2\tsumwx\tsumwx2\tnumEntries\n"
for i in range(len(self.bins)):
s += "{0}\n".format(self.bins[i])
s += "# END YODA_HISTO1D\n"
return s
def Fill(self, x, w):
l = 0
r = len(self.bins) - 1
c = (l + r) // 2
a = self.bins[c].xmin
while r - l > 1:
if x < a:
r = c
else:
l = c
c = (l + r) // 2
a = self.bins[c].xmin
if x > self.bins[r].xmin:
self.oflow.Fill(x, w)
elif x < self.bins[l].xmin:
self.uflow.Fill(x, w)
else:
self.bins[l].Fill(x, w)
self.total.Fill(x, w)
def ScaleW(self, scale):
self.total.ScaleW(scale)
self.uflow.ScaleW(scale)
self.oflow.ScaleW(scale)
for i in range(len(self.bins)):
self.bins[i].ScaleW(scale)
self.scale = scale
def Plot(self, ax=None, *plt_args, **plt_kwargs):
lefts = [b.xmin for b in self.bins]
lefts.append(self.bins[-1].xmax)
heights = [b.w / (lefts[1] - lefts[0]) for b in self.bins]
heights.append(heights[-1])
if ax is None:
fig, ax = plt.subplots()
ax.step(lefts, heights, *plt_args, where='post', **plt_kwargs)
class Scatter2D:
def __init__(self, npoints, xmin, xmax, name="/MC/untitled"):
self.name = name
self.points = []
width = (xmax - xmin) / npoints
for i in range(npoints):
self.points.append(
Point2D(xmin + i * width, xmin + (i + 1) * width))
self.scale = 1.
def __repr__(self):
return str(self)
def __str__(self):
s = "# BEGIN YODA_SCATTER2D {0}\n".format(self.name)
s += "Path={0}\n".format(self.name)
s += "Title=\nType=Histo1D\n"
s += "# xval\txerr-\txerr-\tyval\tyerr-\tyerr+\n"
for i in range(len(self.points)):
s += "{0}\n".format(self.points[i])
s += "# END YODA_SCATTER2D\n"
return s
def ScaleY(self, scale):
for i in range(len(self.points)):
self.points[i].ScaleY(scale)
self.scale = scale
import math as m
from utils.vector import Vec4
class Particle:
def __init__(self, pdgid, momentum, col=[0, 0]):
self.Set(pdgid, momentum, col)
def __repr__(self):
return "{0} {1} {2}".format(self.pid, self.mom, self.col)
def __str__(self):
return "{0} {1} {2}".format(self.pid, self.mom, self.col)
def Set(self, pdgid, momentum, col=[0, 0]):
self.pid = pdgid
self.mom = momentum
self.col = col
def ColorConnected(self, p):
return (self.col[0] > 0 and self.col[0] == p.col[1]) or \
(self.col[1] > 0 and self.col[1] == p.col[0])
def CheckEvent(event):
psum = Vec4()
csum = {}
for p in event:
psum += p.mom
if p.col[0] > 0:
csum[p.col[0]] = csum.get(p.col[0], 0) + 1
if csum[p.col[0]] == 0:
del csum[p.col[0]]
if p.col[1] > 0:
csum[p.col[1]] = csum.get(p.col[1], 0) - 1
if csum[p.col[1]] == 0:
del csum[p.col[1]]
return (m.fabs(psum.E) < 1.e-12 and
m.fabs(psum.px) < 1.e-12 and
m.fabs(psum.py) < 1.e-12 and
m.fabs(psum.pz) < 1.e-12 and
len(csum) == 0)
This diff is collapsed.
import random
from numpy import pi, sqrt, tan, arctan, log, cos, sin
from utils.vector import Vec4
from utils.particle import Particle, CheckEvent
NC = 3.
TR = 1. / 2.
CA = NC
CF = (NC * NC - 1.) / (2. * NC)
class Kernel:
def __init__(self, flavs):
self.flavs = flavs
class Pqq (Kernel):
def Value(self, z, y):
return CF * (2. / (1. - z * (1. - y)) - (1. + z))
def Estimate(self, z):
return CF * 2. / (1. - z)
def Integral(self, zm, zp):
return CF * 2. * log((1. - zm) / (1. - zp))
def GenerateZ(self, zm, zp):
return 1. + (zp - 1.) * pow((1. - zm) / (1. - zp), random.random())
class Pgg (Kernel):
def Value(self, z, y):
return CA / 2. * (2. / (1. - z * (1. - y)) - 2. + z * (1. - z))
def Estimate(self, z):
return CA / (1. - z)
def Integral(self, zm, zp):
return CA * log((1. - zm) / (1. - zp))
def GenerateZ(self, zm, zp):
return 1. + (zp - 1.) * pow((1. - zm) / (1. - zp), random.random())
class Pgq (Kernel):
def Value(self, z, y):
return TR / 2. * (1. - 2. * z * (1. - z))
def Estimate(self, z):
return TR / 2.
def Integral(self, zm, zp):
return TR / 2. * (zp - zm)
def GenerateZ(self, zm, zp):
return zm + (zp - zm) * random.random()
class Shower:
def __init__(self, alpha, t0=1.0):
self.t0 = t0
self.alpha = alpha
self.alphamax = alpha(self.t0)
self.kernels = [Pqq([fl, fl, 21])
for fl in [-5, -4, -3, -2, -1, 1, 2, 3, 4, 5]]
self.kernels += [Pgq([21, fl, -fl]) for fl in [1, 2, 3, 4, 5]]
self.kernels += [Pgg([21, 21, 21])]
def MakeKinematics(self, z, y, phi, pijt, pkt):
Q = pijt + pkt
rkt = sqrt(Q.M2() * y * z * (1. - z))
kt1 = pijt.Cross(pkt)
if kt1.P() < 1.e-6:
kt1 = pijt.Cross(Vec4(0., 1., 0., 0.))
kt1 *= rkt * cos(phi) / kt1.P()
kt2cms = Q.Boost(pijt).Cross(kt1)
kt2cms *= rkt * sin(phi) / kt2cms.P()
kt2 = Q.BoostBack(kt2cms)
pi = z * pijt + (1. - z) * y * pkt + kt1 + kt2
pj = (1. - z) * pijt + z * y * pkt - kt1 - kt2
pk = (1. - y) * pkt
return [pi, pj, pk]
def MakeColors(self, flavs, colij, colk):
self.c += 1
if flavs[0] != 21:
if flavs[0] > 0:
return [[self.c, 0], [colij[0], self.c]]
else:
return [[0, self.c], [self.c, colij[1]]]
else:
if flavs[1] == 21:
if colij[0] == colk[1]:
if colij[1] == colk[0] and random.random() > 0.5:
return [[colij[0], self.c], [self.c, colij[1]]]
return [[self.c, colij[1]], [colij[0], self.c]]
else:
return [[colij[0], self.c], [self.c, colij[1]]]
else:
if flavs[1] > 0:
return [[colij[0], 0], [0, colij[1]]]
else:
return [[0, colij[1]], [colij[0], 0]]
def GeneratePoint(self, event):
while self.t > self.t0:
t = self.t0
for split in event[2:]:
for spect in event[2:]:
if spect == split:
continue
if not split.ColorConnected(spect):
continue
for sf in self.kernels:
if sf.flavs[0] != split.pid:
continue
m2 = (split.mom + spect.mom).M2()
if m2 < 4. * self.t0:
continue
zp = .5 * (1. + sqrt(1. - 4. * self.t0 / m2))
g = self.alphamax / (2. * pi) * \
sf.Integral(1. - zp, zp)
tt = self.t * pow(random.random(), 1. / g)
if tt > t:
t = tt
s = [split, spect, sf, m2, zp]
self.t = t
if t > self.t0:
z = s[2].GenerateZ(1. - s[4], s[4])
y = t / s[3] / z / (1. - z)
if y < 1.:
f = (1. - y) * self.alpha(t) * s[2].Value(z, y)
g = self.alphamax * s[2].Estimate(z)
if f / g > random.random():
phi = 2. * pi * random.random()
moms = self.MakeKinematics(
z, y, phi, s[0].mom, s[1].mom)
cols = self.MakeColors(s[2].flavs, s[0].col, s[1].col)
event.append(Particle(s[2].flavs[2], moms[1], cols[1]))
s[0].Set(s[2].flavs[1], moms[0], cols[0])
s[1].mom = moms[2]
return
def Run(self, event, t):
self.c = 1
self.t = t
while self.t > self.t0:
self.GeneratePoint(event)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment