fixed: cov_ten model

jiangyi15 · Feb 7, 2024 · 6b09a09 · 6b09a09
1 parent 7a2d73a
commit 6b09a09
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 34 deletions.
diff --git a/tf_pwa/amp/cov_ten.py b/tf_pwa/amp/cov_ten.py
@@ -12,7 +12,7 @@
 )
 from tf_pwa.angle import LorentzVector as lv
 from tf_pwa.data import data_shape
-from tf_pwa.particle import _spin_int
+from tf_pwa.particle import BaseDecay, _spin_int
 
 
 class IndexMap:
@@ -99,6 +99,7 @@ def format_idx(index):
         # print(einsum_expr)
         if any(i.dtype is tf.complex128 for i in real_inputs):
             real_inputs = [tf.cast(i, tf.complex128) for i in real_inputs]
+        # print(self, real_inputs)
         # print([i.shape for i in real_inputs])
         ret = tf.einsum(einsum_expr, *real_inputs)
         # print(self.name, ret)
@@ -245,9 +246,9 @@ def __call__(self, inputs, cached=None):
             b = EvalBoost([self.decay.outs[0], self.decay.core], [-1, 1]).call(
                 inputs
             )
-            print("boost matrix")
-            print(b)
-            print("origin t_2^{\\beta\\rho}", ret)
+            # print("boost matrix")
+            # print(b)
+            # print("origin t_2^{\\beta\\rho}", ret)
             b = tf.linalg.inv(b)  #  * np.array([1,-1,-1,-1])
             return tf.einsum("...ab,...bc->...ac", b, ret)
         else:
@@ -385,7 +386,6 @@ def build_decay_einsum(self, ls, idx_map=None):
         return ret
 
     def build_coupling_einsum(self, a, b, c, na, nb, nc, idx_map):
-        # print(a, b, c, na, nb, nc)
         idx_a = [f"{a}_lorentz_{i}" for i in range(na)]
         idx_b = [f"{b}_lorentz_{i}" for i in range(nb)]
         idx_c = [f"{c}_lorentz_{i}" for i in range(nc)]
@@ -421,7 +421,11 @@ def build_coupling_einsum(self, a, b, c, na, nb, nc, idx_map):
             epsilon = EinSum(
                 "epsilon", (idx_a[-1], "delta", idx_b[-1], idx_c[-1])
             )
-            p = EinSum(f"{a}_p", ["delta"])
+            if isinstance(a, str):
+                aa = a.split("->")[0]
+                p = EinSum(f"{aa}_p", ["delta"])
+            else:
+                p = EinSum(f"{a}_p", ["delta"])
             g = [EinSum("g", (i, j)) for i, j in gs]
             ret = EinSum(
                 f"{a}_amp",
@@ -433,9 +437,6 @@ def build_coupling_einsum(self, a, b, c, na, nb, nc, idx_map):
             )
         return ret
 
-    def cal_tensor_t(self, p1, p2, cached_data=None):
-        cached_data = {} if cached_data is None else cached_data
-
     def build_s_einsum(self, decay, l, s, idx_map):
         a = decay.core
         b, c = decay.outs
@@ -510,7 +511,14 @@ def build_wave_function(self, particle, pi):
         if particle.J == 0:
             return tf.ones_like(pi[..., 0:1])
         if particle.J == 1:
-            return wave_function(1, pi)
+            ret = wave_function(1, pi)
+            if len(particle.spins) != 3:
+                ret = tf.gather(
+                    ret,
+                    [_spin_int(i + particle.J) for i in particle.spins],
+                    axis=-1,
+                )
+            return ret
         raise NotImplementedError()
 
 

diff --git a/tf_pwa/tests/config_covten.yml b/tf_pwa/tests/config_covten.yml
@@ -32,7 +32,7 @@ particle:
       Par: 1,
       m0: 2.42,
       g0: 0.03,
-      decay_chain_params: { model: cov_ten },
+      decay_params: { decay_chain_params: { model: cov_ten } },
     }
 
 constrains:

diff --git a/tf_pwa/tests/exp_covten_params.json b/tf_pwa/tests/exp_covten_params.json
@@ -1,26 +1,25 @@
 {
   "A->R_BC.DR_BC->B.C_total_0r": 0.3362381960101486,
-  "A->R_BC.DR_BC->B.C_total_0i": 0.0,
-  "A->R_BC.D_g_ls_1r": 5.302419301961359,
-  "A->R_BC.D_g_ls_1i": -1.6061298241924555,
-  "R_BC->B.C_g_ls_1r": 3.0387310413843807,
-  "R_BC->B.C_g_ls_1i": -12.832877203223928,
-  "R_BC->B.C_g_ls_2r": 4.0568239303382665,
-  "R_BC->B.C_g_ls_2i": -10.286345069408176,
-  "A->R_BD.CR_BD->B.D_total_0r": 0.3045373402432976,
-  "A->R_BD.CR_BD->B.D_total_0i": 1.8968555079221152,
-  "A->R_BD.C_g_ls_1r": 5.411211167748804,
-  "A->R_BD.C_g_ls_1i": -6.990218168324079,
-  "A->R_BD.C_g_ls_2r": 2.783374759721108,
-  "A->R_BD.C_g_ls_2i": 0.5986095451373904,
-  "R_BD->B.D_g_ls_1r": 5.6447152622288,
-  "R_BD->B.D_g_ls_1i": 2.9822710630539095,
-  "A->R_CD.BR_CD->C.D_total_0r": 0.0017709702524416045,
-  "A->R_CD.BR_CD->C.D_total_0i": 8.527976564988037,
-  "A->R_CD.B_g_ls_1r": 2.34367141408338,
-  "A->R_CD.B_g_ls_1i": -12.87738460909358,
-  "A->R_CD.B_g_ls_2r": 3.587995224199312,
-  "A->R_CD.B_g_ls_2i": -10.42326779174883,
-  "R_CD->C.D_g_ls_1r": 24.743983087920917,
-  "R_CD->C.D_g_ls_1i": 1.5638145927928928
+  "A->R_BC.D_g_ls_1r": 4.626907453860775,
+  "A->R_BC.D_g_ls_1i": -1.5193189349197893,
+  "R_BC->B.C_g_ls_1r": 1.562533761390295,
+  "R_BC->B.C_g_ls_1i": -11.583232957494024,
+  "R_BC->B.C_g_ls_2r": 11.673344233144315,
+  "R_BC->B.C_g_ls_2i": -10.356610605516437,
+  "A->R_BD.CR_BD->B.D_total_0r": 0.7412847872707541,
+  "A->R_BD.CR_BD->B.D_total_0i": 2.037586258170449,
+  "A->R_BD.C_g_ls_1r": 3.8730900390701004,
+  "A->R_BD.C_g_ls_1i": -4.991115711610929,
+  "A->R_BD.C_g_ls_2r": 2.9162867284812695,
+  "A->R_BD.C_g_ls_2i": 1.3641617052900632,
+  "R_BD->B.D_g_ls_1r": 3.5887702803674317,
+  "R_BD->B.D_g_ls_1i": 3.4936145812034605,
+  "A->R_CD.BR_CD->C.D_total_0r": 0.0006065596022850162,
+  "A->R_CD.BR_CD->C.D_total_0i": 9.919863136975163,
+  "A->R_CD.B_g_ls_1r": 4.824000012288929,
+  "A->R_CD.B_g_ls_1i": -6.037424964619261,
+  "A->R_CD.B_g_ls_2r": 1.0414997546098455,
+  "A->R_CD.B_g_ls_2i": -7.429464977027848,
+  "R_CD->C.D_g_ls_1r": 34.99794726102977,
+  "R_CD->C.D_g_ls_1i": 1.9261970872406111
 }