Пример #1
0
def test_load():
    with write_temp_file(resonancs_str) as f:
        cs = config_str.format(file_name=f)
        print(cs)
        with write_temp_file(cs) as g:
            config = ConfigLoader(g)
    config.get_amplitude()
Пример #2
0
def test_load():
    with write_temp_file(resonancs_str) as f:
        cs = config_str.format(file_name=f)
        print(cs)
        with write_temp_file(cs) as g:
            config = ConfigLoader(g)
            with open(g) as f:
                data = yaml.full_load(f)
            config2 = ConfigLoader(data)
    config.get_amplitude()
    config2.get_amplitude()
Пример #3
0
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")
Пример #4
0
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.
Пример #5
0
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
Пример #6
0
def test_constrains():
    with write_temp_file(resonancs_str) as f:
        cs = config_str.format(file_name=f)
        print(cs)
        with write_temp_file(cs) as g:
            config = ConfigLoader(g)

    amp = config.get_amplitude()
    config.add_free_var_constraints(amp)
Пример #7
0
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)
Пример #8
0
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")
Пример #9
0
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
Пример #10
0
def test_constrains(gen_toy):
    config = ConfigLoader(f"{this_dir}/config_cfit.yml")
    var_name = "A->R_CD.B_g_ls_1r"
    config.config["constrains"]["init_params"] = {var_name: 1.0}

    @config.register_extra_constrains("init_params")
    def float_var(amp, params=None):
        amp.set_params(params)

    config.register_extra_constrains("init_params2", float_var)

    amp = config.get_amplitude()
    assert amp.get_params()[var_name] == 1.0
Пример #11
0
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)
Пример #12
0
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
Пример #13
0
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)),
    )
Пример #14
0
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)
Пример #15
0
# 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.)}"
    )

# %%
# We can use the amplitude to plot the fit results

amp = config.get_amplitude()
weight = amp(phsp)
partial_weight = amp.partial_weight(phsp)

# %%
# We can plot the data, Hist1D include some plot method base on matplotlib.

data_hist = Hist1D.histogram(data.get_mass("(C, D)"),
                             bins=60,
                             range=(0.25, 1.45))

mass_phsp = phsp.get_mass("(C, D)")
phsp_hist = Hist1D.histogram(mass_phsp,
                             weights=weight,
                             bins=60,
                             range=(0.25, 1.45))
Пример #16
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")