tests: plot with rec

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)

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()

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"]

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

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"