def cal_chi2_config(config,
adapter,
data,
phsp,
data_idx,
bg=None,
data_cut=None):
if data_cut is None:
data_cut = np.array([data_index(data, idx) for idx in data_idx])
amp_weight = config.get_amplitude()(phsp).numpy()
phsp_cut = np.array([data_index(phsp, idx) for idx in data_idx])
phsp_slice = np.concatenate([np.array(phsp_cut**2), [amp_weight]], axis=0)
phsps = adapter.split_data(phsp_slice)
datas = adapter.split_data(data_cut**2)
bound = adapter.get_bounds()
if bg is not None:
bg_weight = config._get_bg_weight(display=False)[0][0]
bg_cut = np.array([data_index(bg, idx) for idx in data_idx])
bgs = adapter.split_data(bg_cut**2)
int_norm = (data_cut.shape[-1] -
bg_cut.shape[-1] * bg_weight) / np.sum(amp_weight)
else:
int_norm = data_cut.shape[-1] / np.sum(amp_weight)
print("int norm:", int_norm)
numbers = []
for i, bnd in enumerate(bound):
min_x, min_y = bnd[0]
max_x, max_y = bnd[1]
ndata = datas[i].shape[-1]
nmc = np.sum(phsps[i][2]) * int_norm
if bg is not None:
nmc += bgs[i].shape[-1] * bg_weight
numbers.append((ndata, nmc))
return cal_chi2(numbers, config.get_ndf())
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")
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")
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 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():
Nbins = 64
config = ConfigLoader("config.yml")
# config = MultiConfig(["config.yml"]).configs[0]
config.set_params("final_params.json")
name = "R_BC"
idx = config.get_data_index("mass", "R_BC")
# idx_costheta = (*config.get_data_index("angle", "DstD/D*"), "beta")
datas, phsps, bgs, _ = config.get_all_data()
amp = config.get_amplitude()
get_data = lambda x: data_index(x, idx).numpy()
# get_data = lambda x: np.cos(data_index(x, idx_costheta).numpy())
plot_mass(amp, datas, bgs, phsps, get_data, name, Nbins)
def cal_chi2(self, read_data=None, bins=[[2, 2]] * 3, mass=["R_BD", "R_CD"]):
if read_data is None:
data_idx = [self.get_data_index("mass", i) for i in mass]
read_data = lambda a: np.array(
[data_index(a, idx)**2 for idx in data_idx])
data, bg, phsp, _ = self.get_all_data()
bg_weight = self._get_bg_weight()[0]
all_data = data_merge(*data)
all_data_cut = read_data(
all_data) # np.array([data_index(a, idx) for idx in data_idx])
adapter = AdaptiveBound(all_data_cut, bins)
numbers = [
self.cal_bins_numbers(adapter, i, j, read_data, k, l)
for i, j, k, l in zip(data, phsp, bg, bg_weight)
]
numbers = np.array(numbers)
numbers = np.sum(numbers, axis=0)
return cal_chi2_o(numbers, self.get_ndf())
m_A, m_B, m_C, m_D = 1.8, 0.18, 0.18, 0.18
p1, p2, p3 = PhaseSpaceGenerator(m_A, [m_B, m_C, m_D]).generate(100000)
# %%
# and the calculate helicity angle from the data
from tf_pwa.cal_angle import cal_angle_from_momentum
data = cal_angle_from_momentum({b: p1, c: p2, d: p3}, decay_group)
# %%
# we can index mass from data as
from tf_pwa.data import data_index
m_bd = data_index(data, ("particle", "(B, D)", "m"))
# m_bc = data_index(data, ("particle", "(B, C)", "m"))
m_cd = data_index(data, ("particle", "(C, D)", "m"))
# %%
# .. note::
# If DecayGroup do not include resonant of (B, C), the data will not include its mass too.
# We can use different DecayGroup for cal_angle and AmplitudeModel
# when they have the same initial and final particle.
#
# The amplitde square is calculated by amplitude model simply as
amp_s2 = amp(data)
# %%
# Now by using matplotlib we can get the Dalitz plot as
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")
def test_plot(init_params):
import matplotlib
matplotlib.rcParams["figure.figsize"] = [10, 10] # for square canvas
matplotlib.rcParams["figure.subplot.left"] = 0.05
matplotlib.rcParams["figure.subplot.bottom"] = 0.05
matplotlib.rcParams["figure.subplot.right"] = 0.95
matplotlib.rcParams["figure.subplot.top"] = 0.95
matplotlib.rcParams["font.size"] = 5
config, amp, data, weight, pw, pw_if = get_data(
init_params=init_params) # init_params=get_params())
m_DsDst = data_index(data,
config.plot_params.get_data_index("mass", "R_BC"))
m_DsK = data_index(data, config.plot_params.get_data_index("mass", "R_BD"))
m_DstK = data_index(data,
config.plot_params.get_data_index("mass", "R_CD"))
val = m_DsDst
sw = np.sum(weight) / weight.shape[0]
print(sw)
# plt.hist(val, weights=[sw]*weight.shape[0], bins=100, label="PHSP")
# plt.hist(val, weights=weight, bins=100, label="data")
label = amp.chains_particle()
label = [str(i[0]) for i in label] #
n_label = len(label)
for i, w in zip(label, pw):
print(i, np.sum(w) / sw)
if n_label % 2 == 0:
a, b = n_label // 2, n_label - 1
else:
a, b = (n_label - 1) // 2, n_label
fig = plt.figure(figsize=(4, 3))
for i, j in enumerate(combinations(range(n_label), 2)):
# ax = plt.subplot(a,b,i+1)
plt.clf()
ax = plt.subplot(1, 1, 1)
ax.hist(val,
weights=pw[j[0]],
bins=100,
histtype="step",
label=label[j[0]])
ax.hist(val,
weights=pw[j[1]],
bins=100,
histtype="step",
label=label[j[1]])
x, y, _ = ax.hist(
val,
weights=pw_if[j],
bins=100,
histtype="step",
label=label[j[0]] + "+" + label[j[1]],
)
ax.hist(
val,
weights=pw_if[j] - pw[j[0]] - pw[j[1]],
bins=100,
histtype="step",
label=label[j[0]] + "x" + label[j[1]],
)
ax.plot(y, np.zeros_like(y))
ax.legend(loc="upper right")
# ax.set_ylim((None, np.max(y) *1.2))
ax.set_title("M(DsD*)")
plt.savefig("fig/fig_{}_{}.pdf".format(*j))