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 main():
config = ConfigLoader("config.yml")
data = config.get_data("data")[0]
phsp = config.get_data("phsp")[0]
bg = config.get_data("bg")[0]
config.set_params("final_params.json")
mbc_idx = config.get_data_index("mass",
"R_BC") # ("particle", "(D, K)", "m")
mbd_idx = config.get_data_index("mass",
"R_BD") # ("particle", "(D0, K, pi)", "m")
data_idx = [mbc_idx, mbd_idx]
data_cut = np.array([data_index(data, idx) for idx in data_idx])
adapter = AdaptiveBound(data_cut**2, [[2, 2], [3, 3], [2, 2]])
bound = adapter.get_bounds()
cal_chi2_config(config,
adapter,
data,
phsp,
data_idx,
bg=bg,
data_cut=data_cut)
draw_dalitz(data_cut, bound)
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():
config = ConfigLoader("config.yml")
decay = config.get_amplitude().decay_group
data = config.get_data("phsp")
ret = []
for i in data:
cached_amp = build_angle_amp_matrix(decay, i)
ret.append(data_to_numpy(cached_amp))
idx = config.get_data_index("angle", "R_BC/B")
ang = data_index(data[0], idx)
np.savez("phsp.npz", ret)
for k, v in ret[0].items():
for i, amp in enumerate(v):
w = np.abs(amp)**2
w = np.sum(np.reshape(w, (amp.shape[0], -1)), axis=-1)
plt.hist(
np.cos(ang["beta"]),
weights=w,
bins=20,
histtype="step",
label="{}: {}".format(k, i),
)
plt.savefig("angle_costheta.png")
def get_data(config_file="config.yml", init_params="init_params.json"):
config = ConfigLoader(config_file)
try:
config.set_params(init_params)
print("using {}".format(init_params))
except Exception as e:
print("using RANDOM parameters")
phsp = config.get_data("phsp")
for i in config.full_decay:
print(i)
for j in i:
print(j.get_ls_list())
print("\n########### initial parameters")
print(json.dumps(config.get_params(), indent=2))
params = config.get_params()
amp = config.get_amplitude()
pw = amp.partial_weight(phsp)
pw_if = amp.partial_weight_interference(phsp)
weight = amp(phsp)
print(weight)
return config, amp, phsp, weight, pw, pw_if
def test_cp_decay():
with open(f"{this_dir}/config_toy.yml") as f:
config_data = yaml.full_load(f)
config_data["decay_chain"] = {"$all": {"is_cp": True}}
config = ConfigLoader(config_data)
amp = config.get_amplitude()
data = config.get_data("data")[0]
amp(data)
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)))
def main():
Nbins = 72
config = ConfigLoader("config.yml")
config.set_params("final_params.json")
error_matrix = np.load("error_matrix.npy")
idx = config.get_data_index("mass", "R_BC")
datas = config.get_data("data")
phsps = config.get_data("phsp")
bgs = config.get_data("bg")
amp = config.get_amplitude()
var = amp.trainable_variables
get_data = lambda x: data_index(x, idx).numpy()
m_all = np.concatenate([get_data(i) for i in datas])
m_min = np.min(m_all) - 0.1
m_max = np.max(m_all) + 0.1
binning = np.linspace(m_min, m_max, Nbins + 1)
get_hist = lambda x, w: Hist1D.histogram(
get_data(x), bins=Nbins, range=(m_min, m_max), weights=w
)
data_hist = [get_hist(i, i.get("weight")) for i in datas]
bg_hist = [get_hist(i, np.abs(i.get("weight"))) for i in bgs]
phsp_hist = []
for dh, bh, phsp in zip(data_hist, bg_hist, phsps):
m_phsp = data_index(phsp, idx).numpy()
y_frac, grads, w_error2 = binning_gradient(
binning, amp, phsp, m_phsp, var
)
error2 = np.einsum("ij,jk,ik->i", grads, error_matrix, grads)
# error parameters and error from integration sample weights
yerr = np.sqrt(error2 + w_error2)
n_fit = dh.get_count() - bh.get_count()
phsp_hist.append(n_fit * Hist1D(binning, y_frac, yerr))
total_data = reduce(operator.add, data_hist)
ax = plt.subplot2grid((4, 1), (0, 0), rowspan=3)
total_data.draw(ax, label="data")
total_data.draw_error(ax)
total_bg = reduce(operator.add, bg_hist)
total_bg.draw_bar(ax, label="back ground", color="grey", alpha=0.5)
total_fit = reduce(operator.add, phsp_hist + bg_hist)
total_fit.draw(ax, label="fit")
total_fit.draw_error(ax)
ax.set_ylim((0, None))
ax.set_xlim((m_min, m_max))
ax.legend()
ax.set_ylabel(f"Events/ {(m_max-m_min)/Nbins:.3f} GeV")
ax2 = plt.subplot2grid((4, 1), (3, 0), rowspan=1)
(total_data - total_fit).draw_pull(ax2)
plt.setp(ax.get_xticklabels(), visible=False)
ax2.set_ylim((-5, 5))
ax2.set_xlim((m_min, m_max))
ax2.axhline(0, c="black", ls="-")
ax2.axhline(-3, c="r", ls="--")
ax2.axhline(3, c="r", ls="--")
ax2.set_xlabel("$M(BC)$/GeV", loc="right")
plt.savefig("fit_full_error.png")