diff --git a/tf_pwa/data_trans/helicity_angle.py b/tf_pwa/data_trans/helicity_angle.py index f0f2c330..1a411d1f 100644 --- a/tf_pwa/data_trans/helicity_angle.py +++ b/tf_pwa/data_trans/helicity_angle.py @@ -164,6 +164,12 @@ def find_variable(self, dat): ms = {topo_map[k]: v for k, v in ms.items()} return ms, costheta, phi + def cal_angle(self, p4): + from tf_pwa.cal_angle import DecayGroup, cal_angle_from_momentum + + decay_group = DecayGroup([self.decay_chain]) + return cal_angle_from_momentum(p4, decay_group) + def mass_linspace(self, name, N): x_min, x_max = self.get_mass_range(name) return np.linspace(x_min + 1e-10, x_max - 1e-10, N) diff --git a/tf_pwa/data_trans/test_helicity_angle.py b/tf_pwa/data_trans/test_helicity_angle.py index 2b1429bb..16a8667c 100644 --- a/tf_pwa/data_trans/test_helicity_angle.py +++ b/tf_pwa/data_trans/test_helicity_angle.py @@ -121,6 +121,12 @@ def test_helicity_var(toy_config): toy_new.get_angle("(C, D)", "C")["alpha"], ) + p_rec = ha.cal_angle(p_new) + p_rec = ha.find_variable(p_rec) + p_rec = ha.build_data(*p_rec) + for k in p_new: + np.allclose(p_rec[k], p_new[k]) + if __name__ == "__main__": test_gen_p() diff --git a/tf_pwa/tests/config_rec.yml b/tf_pwa/tests/config_rec.yml new file mode 100644 index 00000000..dee65cb3 --- /dev/null +++ b/tf_pwa/tests/config_rec.yml @@ -0,0 +1,75 @@ +data: + model: cfit + bg_frac: 0.0909091 + resolution_size: 1 + data_bg_value: "toy_data/data_bg_value.dat" + phsp_bg_value: "toy_data/phsp_bg_value.dat" + data_rec_bg_value: "toy_data/data_bg_value.dat" + phsp_rec_bg_value: "toy_data/phsp_bg_value.dat" + dat_order: [B, C, D] + data: ["toy_data/data.dat"] + data_rec: ["toy_data/data.dat"] + phsp: ["toy_data/PHSP.dat"] + phsp_rec: ["toy_data/PHSP.dat"] + +decay: + A: + - [R_BC, D] + - [R_BD, C] + - [R_CD, B] + R_BC: [B, C] + R_BD: [B, D] + R_CD: [C, D, barrier_factor_norm: True] + +particle: + $top: + A: { J: 1, P: -1, spins: [-1, 1] } + $finals: + B: { J: 1, P: -1 } + C: { J: 1, P: -1 } + D: { J: 0, P: -1 } + R_BC: { J: 1, Par: 1, m0: 4.16, g0: 0.1, model: BWR2 } + R_BD: { J: 1, Par: 1, m0: 2.43, g0: 0.3, model: BW } + R_CD: { J: 1, Par: 1, m0: 2.42, g0: 0.03 } + +constrains: + particle: null + decay: null + var_equal: + - ["A->R_BC.DR_BC->B.C_total_0r", "A->R_BD.CR_BD->B.D_total_0r"] + free_var: + - "R_BC_mass" + +plot: + mass: + R_BC: { display: "$M_{BC}$" } + R_BD: { display: "$M_{BD}$" } + R_CD: { display: "$M_{CD}$" } + angle: + R_BC: + cos(beta): { display: "$cos(\\theta)$" } + aligned_angle: + R_BC/B: + cos(beta): { display: "$cos(\\beta)$" } + 2Dplot: + dalitz_12: + x: m_R_BC**2 + y: m_R_BD**2 + mass_angle: + x: m1 + y: costheta + where: + m1: [mass, R_BC] + costheta: [angle, R_BC, cos(beta)] + xbins: 50 + ybins: 50 + plot_figs: ["data", "sideband_hist", "fitted"] + angle2: + x: a + y: costheta + where: + a: [aligned_angle, R_BC/B, cos(beta)] + costheta: [angle, R_BC, cos(beta)] + xbins: 50 + ybins: 50 + plot_figs: ["data", "sideband_hist", "fitted"] diff --git a/tf_pwa/tests/config_toy3.yml b/tf_pwa/tests/config_toy3.yml index 4a2b4b59..e15fa4bf 100644 --- a/tf_pwa/tests/config_toy3.yml +++ b/tf_pwa/tests/config_toy3.yml @@ -1,8 +1,12 @@ data: dat_order: [B, C, D] + extra_var: + c: + default: 1 data: [["toy_data/data.dat"], ["toy_data/data.dat"]] bg: [["toy_data/bg.dat"], ["toy_data/bg.dat"]] phsp: [["toy_data/PHSP.dat"], ["toy_data/PHSP.dat"]] + phsp_c: ["toy_data/phsp_eff_value.dat", "toy_data/phsp_eff_value.dat"] bg_weight: [0.1, 0.1] random_z: False r_boost: False diff --git a/tf_pwa/tests/test_full.py b/tf_pwa/tests/test_full.py index df410515..3ff2d113 100644 --- a/tf_pwa/tests/test_full.py +++ b/tf_pwa/tests/test_full.py @@ -180,6 +180,16 @@ def test_fit_lazy_call(gen_toy): config.cal_fitfractions() +def test_cfit_resolution(gen_toy): + with open(f"{this_dir}/config_rec.yml") as f: + config_dic = yaml.full_load(f) + config = ConfigLoader(config_dic) + config.set_params(f"{this_dir}/exp_params.json") + fcn = config.get_fcn() + fcn.nll_grad() + config.plot_partial_wave(prefix="toy_data/figure/c3") + + def test_constrains(gen_toy): config = ConfigLoader(f"{this_dir}/config_cfit.yml") var_name = "A->R_CD.B_g_ls_1r"