import numpy as np
import tensorflow as tf
from tf_pwa.amp import get_relative_p
from tf_pwa.angle import LorentzVector as lv
from tf_pwa.data import data_index
[docs]
class HelicityAngle1:
"""simple implement for angle to monmentum trasform"""
def __init__(self, decay_chain):
self.decay_chain = decay_chain
self.par = [i.outs[1] for i in self.decay_chain]
self.par.append(list(self.decay_chain)[-1].outs[0])
[docs]
def generate_p(self, ms, costheta, phi):
msp = [i.get_mass() for i in self.par[:-1]]
return generate_p(ms, msp, costheta, phi)
[docs]
def generate_p2(self, ms, costheta, phi):
all_ms = (
[self.decay_chain[0].core.get_mass()]
+ ms
+ [self.decay_chain[-1].outs[0].get_mass()]
)
return self.generate_p(all_ms, costheta, phi)
[docs]
def generate_p_mass(self, name, m, random=False):
"""generate monmentum with M_name = m"""
m = tf.convert_to_tensor(m, tf.float64)
ms = [
i.core.get_data_mass() if str(i.core) != str(name) else m
for i in self.decay_chain
]
ms.append(self.par[-1].get_mass())
if random:
costheta = [
np.random.random(m.shape) * 2 - 1 for i in self.decay_chain
]
phi = [
np.random.random(m.shape) * 2 * np.pi for i in self.decay_chain
]
else:
costheta = [tf.zeros_like(m) for i in self.decay_chain]
phi = [tf.zeros_like(m) for i in self.decay_chain]
ret = self.generate_p(ms, costheta, phi)
return dict(zip(self.par, ret))
[docs]
def get_phsp_factor(self, name, m):
m = tf.convert_to_tensor(m, tf.float64)
ms = [
i.core.get_mass() if str(i.core) != str(name) else m
for i in self.decay_chain
]
ms.append(self.par[-1].get_mass())
msp = [i.get_mass() for i in self.par[:-1]]
n_decay = len(msp)
ps = [get_relative_p(ms[i], ms[i + 1], msp[i]) for i in range(n_decay)]
ret = tf.ones_like(m)
for i in ps:
ret = ret * tf.cast(i, ret.dtype)
return ret
[docs]
class HelicityAngle:
"""general implement for angle to monmentum trasform"""
def __init__(self, decay_chain):
self.decay_chain = decay_chain
[docs]
def get_all_mass(self, replace_mass):
ms = {}
for i in self.decay_chain:
for j in [i.core] + list(i.outs):
if j not in ms:
if j in replace_mass:
ms[j] = replace_mass[j]
elif str(j) in replace_mass:
ms[j] = tf.convert_to_tensor(
replace_mass[str(j)], tf.float64
)
else:
ms[j] = tf.convert_to_tensor(
j.get_data_mass(), tf.float64
)
return ms
[docs]
def create_ms(self, name, m=None):
if m is None:
if isinstance(name, dict):
ms = self.get_all_mass(name)
m = tf.convert_to_tensor(list(name.values())[0], tf.float64)
else:
raise ValueError("not support input for create_ms")
else:
m = tf.convert_to_tensor(m, tf.float64)
ms = self.get_all_mass({name: m})
return ms, m
[docs]
def generate_p_mass(self, name, m=None, random=False):
"""generate monmentum with M_name = m"""
data = {}
ms, m = self.create_ms(name, m)
for i in self.decay_chain:
data[i] = {}
data[i]["|p|"] = get_relative_p(
data_index(ms, i.core),
data_index(ms, i.outs[0]),
data_index(ms, i.outs[1]),
)
if random:
costheta = np.random.random(m.shape) * 2 - 1
phi = np.random.random(m.shape) * 2 * np.pi
else:
costheta = tf.zeros_like(m)
phi = tf.zeros_like(m)
data[i][i.outs[0]] = {
"angle": {"alpha": phi, "beta": tf.acos(costheta)}
}
ret = create_rotate_p_decay(self.decay_chain, ms, data)
# ret = self.generate_p(ms, costheta, phi)
return ret # dict(zip(self.par, ret))
[docs]
def build_data(self, ms, costheta, phi):
"""generate monmentum with M_name = m"""
data = {}
ms = self.get_all_mass(ms)
for j, i in enumerate(self.decay_chain):
data[i] = {}
data[i]["|p|"] = get_relative_p(
data_index(ms, i.core),
data_index(ms, i.outs[0]),
data_index(ms, i.outs[1]),
)
costheta_i = costheta[j]
phi_i = phi[j]
data[i][i.outs[0]] = {
"angle": {"alpha": phi_i, "beta": tf.acos(costheta_i)}
}
ret = create_rotate_p_decay(self.decay_chain, ms, data)
# ret = self.generate_p(ms, costheta, phi)
return ret # dict(zip(self.par, ret))
[docs]
def get_phsp_factor(self, name, m=None):
ms, _ = self.create_ms(name, m)
return self.eval_phsp_factor(ms)
[docs]
def eval_phsp_factor(self, ms):
ps = []
for i in self.decay_chain:
ps.append(get_relative_p(ms[i.core], ms[i.outs[0]], ms[i.outs[1]]))
ret = tf.ones_like(ms[i.core])
for i in ps:
ret = ret * tf.cast(i, ret.dtype)
return ret
[docs]
def get_mass_range(self, name):
name = str(name)
low_bound = None
high_bound = None
for i in self.decay_chain:
if str(i.core) == name:
low_bound = sum([j.get_data_mass() for j in i.outs])
if name in [str(j) for j in i.outs]:
sum_mass = 0.0
for j in i.outs:
if str(j) != name:
sum_mass = sum_mass + j.get_data_mass()
high_bound = i.core.get_data_mass() - sum_mass
return (low_bound, high_bound)
[docs]
def find_variable(self, dat):
decay_chain = self.decay_chain.standard_topology()
topo_map = decay_chain.topology_map(self.decay_chain)
mi = {decay_chain.top: dat["particle"][decay_chain.top]["m"]}
m2 = {i.outs[1]: dat["particle"][i.outs[1]]["m"] for i in decay_chain}
m1 = {i.outs[0]: dat["particle"][i.outs[0]]["m"] for i in decay_chain}
ang = [
dat["decay"][decay_chain][i][i.outs[0]]["ang"] for i in decay_chain
]
costheta = [tf.cos(i["beta"]) for i in ang]
phi = [i["alpha"] for i in ang]
ms = {**mi, **m1, **m2}
ms = {topo_map[k]: v for k, v in ms.items()}
return ms, costheta, phi
[docs]
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)
[docs]
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)
[docs]
def normal(p):
return p / tf.expand_dims(tf.sqrt(tf.reduce_sum(p**2, axis=-1)), axis=-1)
[docs]
def create_rotate_p(ps, ms, costheta, phi):
px = np.array([[1, 0, 0]])
pz = np.array([[0, 0, 1]])
py = np.array([[0, 1, 0]])
ret = []
for p, m, c, i in zip(ps, ms, costheta, phi):
px_o, py_o, pz_o = px, py, pz
E = tf.sqrt(m * m + p * p)
s = tf.sqrt(1 - c**2)
vx = p * s * tf.cos(i)
vy = p * s * tf.sin(i)
vz = p * c
p_new = (
tf.expand_dims(vx, axis=-1) * px
+ tf.expand_dims(vy, axis=-1) * py
+ tf.expand_dims(vz, axis=-1) * pz
)
E = tf.ones_like(p_new[..., 0]) * E
# print("p2", vx**2+vy**2+vz**2, np.sum(p_new**2, axis=-1), p**2)
ret.append(tf.concat([tf.expand_dims(E, axis=-1), p_new], axis=-1))
pz = normal(p_new)
py = normal(
-px_o * tf.expand_dims(tf.sin(i), axis=-1)
+ py_o * tf.expand_dims(np.cos(i), axis=-1)
)
px = np.cross(py, pz)
return ret
[docs]
def create_rotate_p_decay(decay_chain, mass, data):
px = np.array([[1, 0, 0]])
pz = np.array([[0, 0, 1]])
py = np.array([[0, 1, 0]])
monmentum_in_rest = {}
axis_map = {decay_chain.top: [px, py, pz]}
for _, dec in decay_chain.depth_first():
p, c, i = (
data[dec]["|p|"],
tf.cos(data[dec][dec.outs[0]]["angle"]["beta"]),
data[dec][dec.outs[0]]["angle"]["alpha"],
)
px, py, pz = axis_map[dec.core]
m1 = mass[dec.outs[0]]
E1 = tf.sqrt(m1 * m1 + p * p)
m2 = mass[dec.outs[1]]
E2 = tf.sqrt(m2 * m2 + p * p)
s = tf.sqrt(1 - c**2)
vx = p * s * tf.cos(i)
vy = p * s * tf.sin(i)
vz = p * c
p_new = (
tf.expand_dims(vx, axis=-1) * px
+ tf.expand_dims(vy, axis=-1) * py
+ tf.expand_dims(vz, axis=-1) * pz
)
E1 = tf.ones_like(p_new[..., 0]) * E1
E2 = tf.ones_like(p_new[..., 0]) * E2
p1 = tf.concat([tf.expand_dims(E1, axis=-1), p_new], axis=-1)
p2 = tf.concat([tf.expand_dims(E2, axis=-1), -p_new], axis=-1)
monmentum_in_rest[dec.outs[0]] = p1
monmentum_in_rest[dec.outs[1]] = p2
px_o, py_o, pz_o = px, py, pz
pz = normal(p_new)
py = normal(
-px_o * tf.expand_dims(tf.sin(i), axis=-1)
+ py_o * tf.expand_dims(np.cos(i), axis=-1)
)
px = np.cross(py, pz)
axis_map[dec.outs[0]] = [px, py, pz]
axis_map[dec.outs[1]] = [px, -py, -pz]
monmentum = {}
for _, dec in list(decay_chain.depth_first())[::-1]:
monmentum[dec.core] = {}
for i in dec.outs:
if i in monmentum:
p_boost = monmentum_in_rest[i]
for j in monmentum[i]:
monmentum[dec.core][j] = lv.boost(
monmentum[i][j], lv.boost_vector(p_boost)
)
else:
monmentum[dec.core][i] = monmentum_in_rest[i]
return monmentum[decay_chain.top]
[docs]
def lorentz_neg(pc):
return tf.concat([pc[..., 0:1], -pc[..., 1:]], axis=-1)
[docs]
def generate_p(ms, msp, costheta, phi):
"""
ms(0) -> ms(1) + msp(0), costheta(0), phi(0)
ms(1) -> ms(2) + msp(1), costheta(1), phi(1)
...
ms(n) -> ms(n+1) + msp(n), costheta(n), phi(n)
"""
assert len(ms) == len(msp) + 1
assert len(msp) == len(costheta)
assert len(msp) == len(phi)
n_decay = len(msp)
ms = [tf.convert_to_tensor(i, dtype=tf.float64) for i in ms]
msp = [tf.convert_to_tensor(i, dtype=tf.float64) for i in msp]
costheta = [tf.convert_to_tensor(i, dtype=tf.float64) for i in costheta]
phi = [tf.convert_to_tensor(i, dtype=tf.float64) for i in phi]
ps = [get_relative_p(ms[i], ms[i + 1], msp[i]) for i in range(n_decay)]
p_gen = create_rotate_p(ps, ms[1:], costheta, phi)
p_gen2 = []
for i in range(n_decay):
pa = ps[n_decay - i - 1]
Ea2 = tf.sqrt(msp[n_decay - i - 1] ** 2 + pa**2) + tf.zeros_like(
p_gen[n_decay - i - 1][..., 0]
)
pa2 = tf.concat(
[tf.expand_dims(Ea2, axis=-1), -p_gen[n_decay - i - 1][..., 1:]],
axis=-1,
)
p_gen2.append(pa2)
ret = [p_gen[-1]]
for i in range(n_decay - 1):
boost_p = p_gen[n_decay - i - 2]
pa = lv.boost(p_gen2[i], lv.boost_vector(boost_p))
ret = [lv.boost(j, lv.boost_vector(boost_p)) for j in ret]
ret.append(pa)
ret.append(p_gen2[-1])
return ret[::-1]