def cal_fitfractions(params_file):
config = ConfigLoader("config.yml")
config.set_params(params_file)
params = config.get_params()
config.get_params_error(params)
mcdata = (
config.get_phsp_noeff()
) # use the file of PhaseSpace MC without efficiency indicated in config.yml
fit_frac, err_frac = fit_fractions(
config.get_amplitude(), mcdata, config.inv_he, params
)
print("########## fit fractions:")
fit_frac_string = ""
for i in fit_frac:
if isinstance(i, tuple):
name = "{}x{}".format(*i) # interference term
else:
name = i # fit fraction
fit_frac_string += "{} {}\n".format(
name, error_print(fit_frac[i], err_frac.get(i, None))
)
print(fit_frac_string)
print("########## fit fractions table:")
print_frac_table(
fit_frac_string
) # print the fit-fractions as a 2-D table. The codes below are just to implement the print function.
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 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 fit(config="config.yml", init_params="init_params.json", method="BFGS"):
"""
simple fit script
"""
# load config.yml
config = ConfigLoader(config)
# set initial parameters if have
try:
config.set_params(init_params)
print("using {}".format(init_params))
except Exception as e:
if str(e) != "[Errno 2] No such file or directory: 'init_params.json'":
print(e)
print("\nusing RANDOM parameters", flush=True)
# print("\n########### initial parameters")
# json_print(config.get_params())
# fit
data, phsp, bg, inmc = config.get_all_data()
try:
fit_result = config.fit(batch=65000, method=method)
except KeyboardInterrupt:
config.save_params("break_params.json")
raise
except Exception as e:
print(e)
config.save_params("break_params.json")
raise
json_print(fit_result.params)
fit_result.save_as("final_params.json")
# calculate parameters error
fit_error = config.get_params_error(fit_result, batch=13000)
fit_result.set_error(fit_error)
fit_result.save_as("final_params.json")
pprint(fit_error)
print("\n########## fit results:")
for k, v in config.get_params().items():
print(k, error_print(v, fit_error.get(k, None)))
# plot partial wave distribution
config.plot_partial_wave(fit_result, plot_pull=True)
# calculate fit fractions
phsp_noeff = config.get_phsp_noeff()
fit_frac, err_frac = config.cal_fitfractions({}, phsp_noeff)
print("########## fit fractions")
fit_frac_string = ""
for i in fit_frac:
if isinstance(i, tuple):
name = "{}x{}".format(*i)
else:
name = i
fit_frac_string += "{} {}\n".format(
name, error_print(fit_frac[i], err_frac.get(i, None)))
print(fit_frac_string)
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 generate_toy_from_phspMC(Ndata, mc_file, data_file):
"""Generate toy using PhaseSpace MC from mc_file"""
config = ConfigLoader(f"{this_dir}/config_toy.yml")
config.set_params(f"{this_dir}/gen_params.json")
amp = config.get_amplitude()
data = gen_data(
amp,
Ndata=Ndata,
mcfile=mc_file,
genfile=data_file,
particles=config.get_dat_order(),
)
return data
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 generate_toy_from_phspMC(Ndata, mc_file, data_file):
"""Generate toy using PhaseSpace MC from mc_file"""
# We use ConfigLoader to read the information in the configuration file
config = ConfigLoader("config.yml")
# Set the parameters in the amplitude model
config.set_params("gen_params.json")
amp = config.get_amplitude()
# data is saved in data_file
data = gen_data(
amp,
Ndata=Ndata,
mcfile=mc_file, # input phsase space file
genfile=data_file, # saved toy data file
# use the order in config, the default is ascii order.
particles=config.get_dat_order(),
)
return data
def main():
import argparse
parser = argparse.ArgumentParser(description="calculate fit fractions")
parser.add_argument("-c", "--config", default="config.yml")
parser.add_argument("-i", "--init_params", default="final_params.json")
parser.add_argument("-e", "--error_matrix", default="error_matrix.npy")
results = parser.parse_args()
# load model and parameters and error matrix
config = ConfigLoader(results.config)
config.set_params(results.init_params)
err_matrix = np.load(results.error_matrix)
amp = config.get_amplitude()
phsp = config.get_phsp_noeff() # get_data("phsp")[0]
cal_frac(amp, phsp, err_matrix)
def cal_errors(params_file):
config = ConfigLoader("config.yml")
config.set_params(params_file)
fcn = config.get_fcn()
fcn.model.Amp.vm.rp2xy_all(
) # we can use this to transform all complex parameters to xy-coordinates, since the Hesse errors of xy are more statistically reliable
params = config.get_params()
errors, config.inv_he = cal_hesse_error(
fcn, params, check_posi_def=True, save_npy=True
) # obtain the Hesse errors and the error matrix (inverse Hessian)
print("\n########## fit parameters in XY-coordinates:")
errors = dict(zip(fcn.model.Amp.vm.trainable_vars, errors))
for key, value in config.get_params().items():
print(key, error_print(value, errors.get(key, None)))
print("\n########## correlation matrix:")
print("Matrix index:\n", fcn.model.Amp.vm.trainable_vars)
print("Correlation Coefficients:\n", corr_coef_matrix(config.inv_he)
) # obtain the correlation matrix using the inverse Hessian
def test_cfit(gen_toy):
config = ConfigLoader(f"{this_dir}/config_cfit.yml")
config.set_params(f"{this_dir}/gen_params.json")
fcn = config.get_fcn()
fcn({})
fcn.nll_grad({})
def toy_config(gen_toy):
config = ConfigLoader(f"{this_dir}/config_toy.yml")
config.set_params(f"{this_dir}/exp_params.json")
return config
from tf_pwa.config_loader import ConfigLoader
from tf_pwa.histogram import Hist1D
config = ConfigLoader(yaml.full_load(config_str))
# %%
# We set parameters to a blance value. And we can generate some toy data and calclute the weights
#
input_params = {
"A->R1_a.BR1_a->C.D_total_0r": 6.0,
"A->R1_b.BR1_b->C.D_total_0r": 1.0,
"A->R2.CR2->B.D_total_0r": 2.0,
"A->R3.DR3->B.C_total_0r": 1.0,
}
config.set_params(input_params)
data = config.generate_toy(1000)
phsp = config.generate_phsp(10000)
# You can also fit the data fit to the data
fit_result = config.fit([data], [phsp])
err = config.get_params_error(fit_result, [data], [phsp])
# %%
# we can see that thre fit results consistant with inputs, the first one is fixed.
for var in input_params:
print(
f"in: {input_params[var]} => out: {fit_result.params[var]} +/- {err.get(var, 0.)}"
)
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")