from tf_pwa.tensorflow_wrapper import tf
from .core import Particle, Variable, register_particle
[docs]def cal_monentum(m, ma, mb):
mabp = ma + mb
mabm = ma - mb
s = m * m
p2 = (s - mabp * mabp) * (s - mabm * mabm) / 4 / s
zeros = tf.zeros_like(s)
p_p = tf.complex(tf.sqrt(tf.abs(p2)), zeros)
p_m = tf.complex(zeros, tf.sqrt(tf.abs(p2)))
return tf.where(p2 > 0, p_p, p_m)
[docs]@register_particle("Flatte")
class ParticleFlatte(Particle):
def __init__(self, *args, mass_list=None, **kwargs):
super().__init__(*args, **kwargs)
if mass_list is None:
raise ValueError("required mass_list: [[a, b], [mc, md]]")
self.mass_list = mass_list
self.g_value = []
self.float_list = list(kwargs.get("float", []))
[docs] def init_params(self):
self.d = 3.0
if self.mass is None:
self.mass = self.add_var("mass", fix=True)
# print("$$$$$",self.mass)
else:
if not isinstance(self.mass, Variable):
if "m" in self.float_list:
self.mass = self.add_var(
"mass", value=self.mass, fix=False
)
else:
self.mass = self.add_var("mass", value=self.mass, fix=True)
self.g_value = []
for i, mab in enumerate(self.mass_list):
name = f"g_{i}"
if hasattr(self, name):
if name in self.float_list:
self.g_value.append(
self.add_var(
f"g_{i}", value=getattr(self, name), fix=False
)
)
else:
self.g_value.append(
self.add_var(
f"g_{i}", value=getattr(self, name), fix=True
)
)
else:
self.g_value.append(self.add_var(f"g_{i}"))
def __call__(self, m):
return self.get_amp({"m": m})
[docs] def get_amp(self, *args, **kwargs):
m = args[0]["m"]
mass = self.get_mass()
zeros = tf.zeros_like(m)
delta_s = mass * mass - m * m
m_c = mass / m
rhos = []
for i, mab in enumerate(self.mass_list):
ma, mb = mab
pi = cal_monentum(m, ma, mb)
# print(pi)
m_rho_i = pi * tf.complex(zeros, self.g_value[i]() * m_c)
rhos.append(m_rho_i)
rho = sum(rhos)
re = delta_s + tf.math.real(rho)
im = tf.math.imag(rho)
d = re * re + im * im
ret = tf.complex(re / d, -im / d)
return ret