diff --git a/tf_pwa/amp/cov_ten.py b/tf_pwa/amp/cov_ten.py index 4054286..ea9f2dc 100644 --- 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 index b21b3ee..1517668 100644 --- 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 index 23301da..6377d47 100644 --- 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 }