def main():
sigma = 0.005
r_name = "R_BC"
config = ConfigLoader("config.yml")
sample_N = config.resolution_size
decays = config.get_decay(False)
decay_chain = decays.get_decay_chain(r_name)
data = config.get_data("data_origin")[0]
phsp = config.get_data("phsp_plot")[0]
dat_order = config.get_dat_order()
generate = lambda x: gauss_random(x, decay_chain, r_name, sigma, sample_N,
dat_order)
pi = generate(data)
np.savetxt("data/data.dat", pi)
np.savetxt("data/data_weight.dat",
np.ones((pi.shape[0] // len(dat_order), )))
pi = generate(phsp)
np.save("data/phsp_re.npy", pi)
np.savetxt("data/phsp_re_weight.dat",
np.ones((pi.shape[0] // len(dat_order), )))
def plot_all(
res="MI(1+)S",
config_file="config.yml",
params="final_params.json",
prefix="figure/",
):
"""plot all figure"""
config = ConfigLoader(config_file)
config.set_params(params)
particle = config.get_decay().get_particle(res)
mi, r, phi_i, r_e, phi_e = load_params(config_file, params, res)
x, y, x_e, y_e = trans_r2xy(r, phi_i, r_e, phi_e)
m = np.linspace(mi[0], mi[-1], 1000)
M_Kpm = 0.49368
M_Dpm = 1.86961
M_Dstar0 = 2.00685
M_Bpm = 5.27926
# x_new = interp1d(xi, x, "cubic")(m)
# y_new = interp1d(xi, y, "cubic")(m)
rm_new = particle.interp(m).numpy()
x_new, y_new = rm_new.real, rm_new.imag
pq = dalitz_weight(m * m, M_Bpm, M_Dstar0, M_Dpm, M_Kpm)
pq_i = dalitz_weight(mi * mi, M_Bpm, M_Dstar0, M_Dpm, M_Kpm)
phi = np.arctan2(y_new, x_new)
r2 = x_new * x_new + y_new * y_new
plot_phi(f"{prefix}phi.png", m, phi, mi, np.arctan2(y, x))
plot_x_y(
f"{prefix}r2.png",
m,
r2,
mi,
r * r,
"mass",
"$|R(m)|^2$",
ylim=(0, None),
)
plot_x_y(f"{prefix}x_y.png", x_new, y_new, x, y, "real R(m)", "imag R(m)")
plot_x_y_err(f"{prefix}x_y_err.png", x[1:-1], y[1:-1], x_e[1:-1],
y_e[1:-1])
plot_x_y(
f"{prefix}r2_pq.png",
m,
r2 * pq,
mi,
r * r * pq_i,
"mass",
"$|R(m)|^2 p \cdot q$",
ylim=(0, None),
)
plot3d_m_x_y(f"{prefix}m_r.gif", m, x_new, y_new)
def main():
config = ConfigLoader("config.yml")
config.set_params("final_params.json")
amp = config.get_amplitude()
data = config.get_data("data_origin")[0]
phsp = config.get_data("phsp_plot")[0]
phsp_re = config.get_data("phsp_plot_re")[0]
print("data loaded")
amps = amp(phsp_re)
pw = amp.partial_weight(phsp_re)
re_weight = phsp_re["weight"]
re_size = config.resolution_size
amps = sum_resolution(amps, re_weight, re_size)
pw = [sum_resolution(i, re_weight, re_size) for i in pw]
m_idx = config.get_data_index("mass", "R_BC")
m_phsp = data_index(phsp, m_idx).numpy()
m_data = data_index(data, m_idx).numpy()
m_min, m_max = np.min(m_phsp), np.max(m_phsp)
scale = m_data.shape[0] / np.sum(amps)
get_hist = lambda m, w: Hist1D.histogram(
m, weights=w, range=(m_min, m_max), bins=100)
data_hist = get_hist(m_data, None)
phsp_hist = get_hist(m_phsp, scale * amps)
pw_hist = []
for i in pw:
pw_hist.append(get_hist(m_phsp, scale * i))
ax2 = plt.subplot2grid((4, 1), (3, 0), rowspan=1)
ax = plt.subplot2grid((4, 1), (0, 0), rowspan=3, sharex=ax2)
data_hist.draw_error(ax, label="data")
phsp_hist.draw(ax, label="fit")
for i, j in zip(pw_hist, config.get_decay()):
i.draw_kde(ax, label=str(j.inner[0]))
(data_hist - phsp_hist).draw_pull(ax2)
ax.set_ylim((1, None))
ax.legend()
ax.set_yscale("log")
ax.set_ylabel("Events/{:.1f} MeV".format((m_max - m_min) * 10))
ax2.set_xlabel("M( R_BC )")
ax2.set_ylabel("pull")
ax2.set_xlim((1.3, 1.7))
ax2.set_ylim((-5, 5))
plt.setp(ax.get_xticklabels(), visible=False)
plt.savefig("m_R_BC_fit.png")
def main():
sigma = 0.005
config = ConfigLoader("config.yml")
decay = config.get_decay()
m0 = decay.top.get_mass()
m1, m2, m3 = [i.get_mass() for i in decay.outs]
print("mass: ", m0, " -> ", m1, m2, m3)
phsp = PhaseSpaceGenerator(m0, [m1, m2, m3])
p1, p2, p3 = phsp.generate(100000)
angle = cal_angle_from_momentum({"B": p1, "C": p2, "D": p3}, decay)
amp = config.get_amplitude()
m_idx = config.get_data_index("mass", "R_BC")
m_BC = data_index(angle, m_idx)
R_BC = decay.get_particle("R_BC1")
m_R, g_R = R_BC.get_mass(), R_BC.get_width()
# import matplotlib.pyplot as plt
# x = np.linspace(1.3, 1.7, 1000)
# amp = R_BC.get_amp({"m": x}).numpy()
# plt.plot(x, np.abs(amp)**2)
# plt.show()
print("mass: ", m_R, "width: ", g_R)
amp_s2 = resolution_bw(m_BC, m_R, g_R, sigma, m1 + m2, m0 - m3)
print("|A|*R: ", amp_s2)
cut_data = simple_selection(angle, amp_s2)
ps = [data_index(cut_data, ("particle", i, "p")).numpy() for i in "BCD"]
np.savetxt("data/data_origin.dat",
np.transpose(ps, (1, 0, 2)).reshape((-1, 4)))
p1, p2, p3 = phsp.generate(100000)
np.savetxt("data/phsp.dat",
np.transpose([p1, p2, p3], (1, 0, 2)).reshape((-1, 4)))
p1, p2, p3 = phsp.generate(50000)
np.savetxt(
"data/phsp_plot.dat",
np.transpose([p1, p2, p3], (1, 0, 2)).reshape((-1, 4)),
)
def main():
sym.init_printing()
config = ConfigLoader("config.yml")
decay_group = config.get_decay()
decay_chain = list(decay_group)[0]
print(decay_chain)
angle = get_angle_distrubution(decay_chain)
ret = []
for k, v in angle.items():
for k2, v2 in v.items():
f_theta = simplify(v2 * v2.conjugate())
print("ls: ", k)
print(" lambda:", k2)
print(" :", f_theta)
ret.append(f_theta)
# break
f_theta = ret[0]
plot_theta("costheta.png", f_theta, "theta1")
plot_theta("costheta2.png", f_theta, "theta2")
plot_theta("phi2.png", f_theta, "phi2")
def main():
sigma = 0.005
sigma_delta = 5
r_name = "R_BC"
config = ConfigLoader("config.yml")
sample_N = config.resolution_size
decays = config.get_decay(False)
decay_chain = decays.get_decay_chain(r_name)
data = config.get_data("data_origin")[0]
pi, total_weights = gauss_sample(data, decay_chain, "R_BC", sigma,
sample_N, config.get_dat_order())
np.savetxt("data/data.dat", pi.reshape((-1, 4)))
np.savetxt("data/data_weight.dat", np.reshape(total_weights, (-1, )))
data = config.get_data("phsp_plot")[0]
pi, total_weights = gauss_sample(data, decay_chain, "R_BC", sigma,
sample_N, config.get_dat_order())
np.save("data/phsp_re.npy", pi.reshape((-1, 4)))
np.savetxt("data/phsp_re_weight.dat", np.reshape(total_weights, (-1, )))
def main():
sigma = 0.005
r_name = "DstD"
config = ConfigLoader("config_data.yml")
decays = config.get_decay(False)
decay_chain = decays.get_decay_chain(r_name)
data = config.get_data("data")[0]
angle = cal_helicity_angle(data["particle"],
decay_chain.standard_topology())
decay_chain.standard_topology()
tp_map = decay_chain.topology_map()
r_particle = tp_map[get_particle(r_name)]
mass = {}
for i in data["particle"]:
mi = data["particle"][i]["m"]
if i == r_particle:
mi = mi + tf.random.normal(mi.shape, 0, sigma, dtype=mi.dtype)
mass[i] = mi
mask = True
p4_all = {}
for i in decay_chain:
phi = angle[tp_map[i]][tp_map[i.outs[0]]]["ang"]["alpha"]
theta = angle[tp_map[i]][tp_map[i.outs[0]]]["ang"]["beta"]
m0 = mass[tp_map[i.core]]
m1 = mass[tp_map[i.outs[0]]]
m2 = mass[tp_map[i.outs[1]]]
mask = mask & (m0 >= m1 + m2)
p_square = get_relative_p2(m0, m1, m2)
p = tf.sqrt(tf.where(p_square > 0, p_square, 0))
pz = p * tf.cos(theta)
px = p * tf.sin(theta) * tf.cos(phi)
py = p * tf.sin(theta) * tf.sin(phi)
E1 = tf.sqrt(m1 * m1 + p * p)
E2 = tf.sqrt(m2 * m2 + p * p)
p1 = tf.stack([E1, px, py, pz], axis=-1)
p2 = tf.stack([E2, -px, -py, -pz], axis=-1)
p4_all[i.outs[0]] = p1
p4_all[i.outs[1]] = p2
core_boost = {}
for i in decay_chain:
if i.core != decay_chain.top:
core_boost[i.outs[0]] = i.core
core_boost[i.outs[1]] = i.core
ret = {}
for i in decay_chain.outs:
tmp = i
ret[i] = p4_all[i]
while tmp in core_boost:
tmp = core_boost[tmp]
# print(i, tmp)
ret[i] = lv.rest_vector(lv.neg(p4_all[tmp]), ret[i])
ret2 = {}
mask = tf.expand_dims(mask, -1)
for i in ret:
ret2[i] = tf.where(mask, ret[i], data["particle"][tp_map[i]]["p"])
# print(ret)
# print({i: data["particle"][tp_map[i]]["p"] for i in decay_chain.outs})
pi = np.stack([ret2[i] for i in config.get_dat_order()], axis=1)
np.savetxt("data_gauss.dat", pi.reshape((-1, 4)))
mass_phsp = phsp.get_mass("(C, D)")
phsp_hist = Hist1D.histogram(mass_phsp,
weights=weight,
bins=60,
range=(0.25, 1.45))
scale = phsp_hist.scale_to(data_hist)
pw_hist = []
for w in partial_weight:
# here we used more bins for a smooth plot
hist = Hist1D.histogram(mass_phsp,
weights=w,
bins=60 * 2,
range=(0.25, 1.45))
pw_hist.append(scale * hist * 2)
# %%
# Then we can plot the histogram into matplotlib
for hist, dec in zip(pw_hist, config.get_decay()):
hist.draw_kde(label=str(dec))
phsp_hist.draw(label="total amplitude")
data_hist.draw_error(label="toy data", color="black")
plt.legend()
plt.ylim((0, None))
plt.xlabel("M(CD)/ GeV")
plt.ylabel("Events/ 20 MeV")
plt.show()
def main():
config = ConfigLoader("config.yml")
for i, dec in enumerate(config.get_decay()):
draw_decay_struct(dec,
filename="figure/fig_{}".format(i),
format="png")