def test_spline_c_idx():
p = get_particle(
"a",
model="spline_c_idx",
min_m=1.0,
max_m=3.0,
interp_N=20,
with_bound=True,
)
with variable_scope() as vm:
p.init_params()
for i in range(20):
vm.set(f"a_point_{i}r", 1.0)
vm.set(f"a_point_{i}i", 0.0)
amp = p(np.array([1.0, 3.0, 2.0]))
assert np.allclose(amp, 1.0)
p = get_particle(
"a",
model="spline_c_idx",
polar=False,
points=[1.0, 2.0, 3.0, 4.0, 5.0],
)
with variable_scope() as vm:
p.init_params()
for i in range(3):
vm.set(f"a_point_{i}r", i)
vm.set(f"a_point_{i}i", i)
amp = p(np.array([1.0, 3.0, 2.0, 4.0]))
assert np.allclose(amp, [0.0, 1 + 1j, 0.0, 2 + 2j])
def get_angle_vars(self):
ang = self.config.get("angle", {})
for k, i in ang.items():
names = k.split("/")
name = names[0]
number_decay = True
if len(names) > 1:
try:
count = int(names[-1])
except ValueError:
number_decay = False
else:
count = 0
if number_decay:
decay_chain, decay = None, None
part = self.re_map.get(get_particle(name), get_particle(name))
for decs in self.decay_struct:
for dec in decs:
if dec.core == get_particle(name):
decay = dec.core.decay[count]
for j in self.decay_struct:
if decay in j:
decay_chain = j.standard_topology()
decay = self.re_map.get(decay, decay)
part = decay.outs[0]
else:
_, decay_chain, decay, part, _ = self.get_data_index(
"angle", k)
for j, v in i.items():
display = v.get("display", j)
upper_ylim = v.get("upper_ylim", None)
theta = j
trans = lambda x: x
if "cos" in j:
theta = j[4:-1]
trans = np.cos
bins = v.get("bins", self.defaults_config.get("bins", 50))
xrange = v.get("range", None)
legend = v.get("legend",
self.defaults_config.get("legend", False))
yscale = v.get("yscale",
self.defaults_config.get("yscale", "linear"))
yield {
"name": validate_file_name(k + "_" + j),
"display": display,
"upper_ylim": upper_ylim,
"idx": ("decay", decay_chain, decay, part, "ang", theta),
"trans": trans,
"bins": bins,
"range": xrange,
"legend": legend,
"yscale": yscale,
}
def get_dat_order(self, standard=False):
order = self.dic.get("dat_order", None)
if order is None:
order = list(self.decay_struct.outs)
else:
order = [get_particle(str(i)) for i in order]
if not standard:
return order
re_map = self.decay_struct.get_chains_map()
def particle_item():
for j in re_map:
for k, v in j.items():
for s, l in v.items():
yield s, l
new_order = []
for i in order:
for s, l in particle_item():
if str(l) == str(i):
new_order.append(s)
break
else:
new_order.append(i)
return new_order
def possible_jp(decay_chain, max_J=2):
"""get possible resonances jp of J <= max_J"""
ret = []
A = decay_chain.top
outs = decay_chain.outs
jp_list = [jp_seq(max_J) for _ in decay_chain.inner]
ret = {}
for jps in itertools.product(*jp_list):
res_map = {}
for jp, res in zip(jps, decay_chain.inner):
j, p = jp
res_n = get_particle("{}_{}{}".format(res, j, p), J=j, P=p)
res_map[res] = res_n
possible_ls = []
for i in decay_chain:
core = res_map.get(i.core, i.core)
outs = [res_map.get(j, j) for j in i.outs]
dec = get_decay(core, outs, p_break=i.p_break, disable=True)
# print(core.J, core.P, outs[0].J, outs[0].P, outs[1].J, outs[1].P)
# print(dec.get_ls_list())
possible_ls.append(dec.get_ls_list())
if min([len(i) for i in possible_ls]) > 0:
ret[jps] = dict(zip(decay_chain, possible_ls))
return ret
def load_data(self,
files,
weights=None,
weights_sign=1,
charge=None) -> dict:
# print(files, weights)
if files is None:
return None
order = self.get_dat_order()
p_list = self.get_particle_p()
center_mass = self.dic.get("center_mass", True)
r_boost = self.dic.get("r_boost", False)
random_z = self.dic.get("random_z", False)
npz_data = np.load(files)
p = {
get_particle(str(v)): npz_data[str(k)]
for k, v in zip(p_list, order)
}
data = cal_angle_from_momentum(
p,
self.decay_struct,
center_mass=center_mass,
r_boost=r_boost,
random_z=random_z,
)
if "weight" in npz_data:
data["weight"] = npz_data["weight"]
if "charge_conjugation" in npz_data:
data["charge_conjugation"] = npz_data["charge_conjugation"]
else:
data["charge_conjugation"] = np.ones((data_shape(data), ))
return data
def get_data_index(self, sub, name):
dec = self.decay_struct.topology_structure()
if sub == "mass":
p = get_particle(name)
return "particle", self.re_map.get(p, p), "m"
if sub == "p":
p = get_particle(name)
return "particle", self.re_map.get(p, p), "p"
if sub == "angle":
name_i = name.split("/")
de_i = self.decay_struct.get_decay_chain(name_i)
p = get_particle(name_i[-1])
for i in de_i:
if p in i.outs:
de = i
break
else:
raise IndexError("not found such decay {}".format(name))
return (
"decay",
de_i.standard_topology(),
self.re_map.get(de, de),
self.re_map.get(p, p),
"ang",
)
if sub == "aligned_angle":
name_i = name.split("/")
de_i = self.decay_struct.get_decay_chain(name_i)
p = get_particle(name_i[-1])
for i in de_i:
if p in i.outs:
de = i
break
else:
raise IndexError("not found such decay {}".format(name))
return (
"decay",
de_i.standard_topology(),
self.re_map.get(de, de),
self.re_map.get(p, p),
"aligned_angle",
)
raise ValueError("unknown sub {}".format(sub))
def test_hist_idx():
p = get_particle("a", model="hist_idx", min_m=1.0, max_m=3.0, interp_N=20)
with variable_scope() as vm:
p.init_params()
vm.set("a_point_0r", 1.0)
vm.set("a_point_0i", 0.0)
vm.set("a_point_1r", 1.0)
vm.set("a_point_1i", 0.0)
amp = p(np.array([1.0, 1.001, 2.0, 3.0, 4.0]))
assert np.allclose(amp[0:2], 1.0)
def add_particle(name):
if name in particle_set:
return particle_set[name]
params = particle_params.get(name, {})
params = self.rename_params(params)
set_min_max(params, "mass", "m_min", "m_max")
set_min_max(params, "width", "g_min", "g_max")
part = get_particle(name, **params)
particle_set[name] = part
return part
def get_mass_vars(self):
mass = self.config.get("mass", {})
x = sy.symbols("x")
for k, v in mass.items():
display = v.get("display", "M({})".format(k))
upper_ylim = v.get("upper_ylim", None)
xrange = v.get("range", None)
trans = v.get("trans", None)
if trans is None:
trans = lambda x: x
else:
trans = sy.sympify(trans)
trans = sy.lambdify(x, trans, modules="numpy")
units = v.get("units", "GeV")
bins = v.get("bins", self.defaults_config.get("bins", 50))
legend = v.get("legend", self.defaults_config.get("legend", True))
yscale = v.get("yscale",
self.defaults_config.get("yscale", "linear"))
yield {
"name":
"m_" + k,
"display":
display,
"upper_ylim":
upper_ylim,
"idx": (
"particle",
self.re_map.get(get_particle(k), get_particle(k)),
"m",
),
"legend":
legend,
"range":
xrange,
"bins":
bins,
"trans":
trans,
"units":
units,
"yscale":
yscale,
}
def main():
a = get_particle("A", J=0, P=-1)
b = get_particle("B", J=1, P=-1)
c = get_particle("C", J=0, P=-1)
d = get_particle("D", J=0, P=-1)
r = get_particle("R")
dec1 = get_decay(a, [r, d], p_break=True)
dec2 = get_decay(r, [b, c])
decs = DecayChain([dec1, dec2])
# chains = DecayChain.from_particles(a, [b, c, d])
# for i in chains:
jp_list = possible_jp(decs)
for jps in jp_list:
sign = lambda x: "+" if x == 1 else "-"
jp_display = ["{}{}".format(j, sign(p)) for j, p in jps]
print("Resonances JP: ", dict(zip(decs.inner, jp_display)))
print(" Possible (l, s): ", jp_list[jps])
def test_a():
a = get_particle(
"a",
mass_list=[2.0],
width_list=[0.03],
m1=1.1,
m2=0.5,
bw_l=1,
model="KMatrixSingleChannel",
)
b = get_particle("b", J=1, P=-1, mass=2.0, width=0.03, model="BWR")
c = get_particle("c", J=0, P=-1, mass=1.1, width=0.03, model="one")
d = get_particle("d", J=0, P=-1, mass=0.5, width=0.03, model="one")
dec = get_decay(b, [c, d])
a.init_params()
b.init_params()
m = np.array([1.9, 2.0, 2.1])
q = get_relative_p(m, 1.1, 0.5)
q0 = get_relative_p(2.0, 1.1, 0.5)
ay = a.get_amp({"m": m})
by = b.get_amp({"m": m}, {"|q|": q, "|q0|": q0})
print(ay / by)
check_frac(ay.numpy(), by.numpy())
def main():
sigma = 0.005
r_name = "DstD"
config = ConfigLoader("config_data.yml")
decays = config.get_decay(False)
decay_chain = decays.get_decay_chain(r_name)
data = config.get_data("data")[0]
angle = cal_helicity_angle(data["particle"],
decay_chain.standard_topology())
decay_chain.standard_topology()
tp_map = decay_chain.topology_map()
r_particle = tp_map[get_particle(r_name)]
mass = {}
for i in data["particle"]:
mi = data["particle"][i]["m"]
if i == r_particle:
mi = mi + tf.random.normal(mi.shape, 0, sigma, dtype=mi.dtype)
mass[i] = mi
mask = True
p4_all = {}
for i in decay_chain:
phi = angle[tp_map[i]][tp_map[i.outs[0]]]["ang"]["alpha"]
theta = angle[tp_map[i]][tp_map[i.outs[0]]]["ang"]["beta"]
m0 = mass[tp_map[i.core]]
m1 = mass[tp_map[i.outs[0]]]
m2 = mass[tp_map[i.outs[1]]]
mask = mask & (m0 >= m1 + m2)
p_square = get_relative_p2(m0, m1, m2)
p = tf.sqrt(tf.where(p_square > 0, p_square, 0))
pz = p * tf.cos(theta)
px = p * tf.sin(theta) * tf.cos(phi)
py = p * tf.sin(theta) * tf.sin(phi)
E1 = tf.sqrt(m1 * m1 + p * p)
E2 = tf.sqrt(m2 * m2 + p * p)
p1 = tf.stack([E1, px, py, pz], axis=-1)
p2 = tf.stack([E2, -px, -py, -pz], axis=-1)
p4_all[i.outs[0]] = p1
p4_all[i.outs[1]] = p2
core_boost = {}
for i in decay_chain:
if i.core != decay_chain.top:
core_boost[i.outs[0]] = i.core
core_boost[i.outs[1]] = i.core
ret = {}
for i in decay_chain.outs:
tmp = i
ret[i] = p4_all[i]
while tmp in core_boost:
tmp = core_boost[tmp]
# print(i, tmp)
ret[i] = lv.rest_vector(lv.neg(p4_all[tmp]), ret[i])
ret2 = {}
mask = tf.expand_dims(mask, -1)
for i in ret:
ret2[i] = tf.where(mask, ret[i], data["particle"][tp_map[i]]["p"])
# print(ret)
# print({i: data["particle"][tp_map[i]]["p"] for i in decay_chain.outs})
pi = np.stack([ret2[i] for i in config.get_dat_order()], axis=1)
np.savetxt("data_gauss.dat", pi.reshape((-1, 4)))
},
"R1": {
"J": 0,
"P": 1,
"mass": 1.0,
"width": 0.07
},
"R2": {
"J": 1,
"P": -1,
"mass": 1.225,
"width": 0.08
},
}
a, b, c, d = [get_particle(i, J=0, P=-1) for i in "ABCD"]
r1, r2, r3 = [get_particle(i, **resonances[i]) for i in resonances.keys()]
decay_group = DecayGroup([
DecayChain([get_decay(a, [r1, c]),
get_decay(r1, [b, d])]),
DecayChain([get_decay(a, [r2, b]),
get_decay(r2, [c, d])]),
DecayChain([get_decay(a, [r3, b]),
get_decay(r3, [c, d])]),
])
# %%
# The above parts can be represented as config.yml used by ConfigLoader.
#
# We can get AmplitudeModel form decay_group and a optional Variables Managerager.
def gauss_sample(data, decay_chain, r_name, sigma, sample_N, dat_order):
sigma_delta = 5
def gauss(delta_x):
return tf.exp(-(delta_x**2) / (2 * sigma**2))
angle = cal_helicity_angle(data["particle"],
decay_chain.standard_topology())
decay_chain.standard_topology()
tp_map = decay_chain.topology_map()
r_particle = tp_map[get_particle(r_name)]
for i in decay_chain.standard_topology():
if i.core == r_particle:
m_min = sum(data["particle"][j]["m"] for j in i.outs)
print(i.outs, r_particle)
if any(r_particle == j for j in i.outs):
m_max = (data["particle"][i.core]["m"] -
sum(data["particle"][j]["m"]
for j in i.outs) + data["particle"][r_particle]["m"])
print("min, max: ", m_min, m_max)
mass = {}
weights = []
for i in data["particle"]:
mi = data["particle"][i]["m"]
if i == r_particle:
delta_min = tf.where(
m_min - mi > -sigma_delta * sigma,
m_min - mi,
-sigma_delta * sigma,
)
delta_max = tf.where(
m_max - mi > sigma_delta * sigma,
sigma_delta * sigma,
m_max - mi,
)
delta_m = (delta_max - delta_min) / (sample_N + 1)
print("delta_min:", delta_min)
min_m = mi + delta_min + delta_m / 2
mi_s = []
for j in range(sample_N):
mi_s_i = min_m + delta_m * j
mi_s.append(mi_s_i)
weights.append(gauss(mi_s_i - mi))
mass[i] = tf.stack(mi_s)
else:
mass[i] = mi[None, :]
# print(mass[r_particle], np.mean(mass[r_particle]))
weights = tf.stack(weights)
weights = weights / tf.reduce_sum(weights, axis=0)
data_weights = data.get("weight", tf.ones_like(weights))
total_weights = weights * data_weights
print({k: v.shape for k, v in mass.items()})
mask = True
p4_all = {}
for i in decay_chain:
phi = angle[tp_map[i]][tp_map[i.outs[0]]]["ang"]["alpha"]
theta = angle[tp_map[i]][tp_map[i.outs[0]]]["ang"]["beta"]
m0 = mass[tp_map[i.core]]
m1 = mass[tp_map[i.outs[0]]]
m2 = mass[tp_map[i.outs[1]]]
p_square = get_relative_p2(m0, m1, m2)
print(m0.shape, m1.shape, m2.shape, p_square.shape)
p = tf.sqrt(tf.where(p_square > 0, p_square, 0))
pz = p * tf.cos(theta)
px = p * tf.sin(theta) * tf.cos(phi)
py = p * tf.sin(theta) * tf.sin(phi)
E1 = tf.sqrt(m1 * m1 + p * p)
E2 = tf.sqrt(m2 * m2 + p * p)
p1 = tf.stack([E1, px, py, pz], axis=-1)
p2 = tf.stack([E2, -px, -py, -pz], axis=-1)
p4_all[i.outs[0]] = p1
p4_all[i.outs[1]] = p2
print("p shape", {k: v.shape for k, v in p4_all.items()})
core_boost = {}
for i in decay_chain:
if i.core != decay_chain.top:
core_boost[i.outs[0]] = i.core
core_boost[i.outs[1]] = i.core
ret = {}
for i in decay_chain.outs:
tmp = i
ret[i] = p4_all[i]
while tmp in core_boost:
tmp = core_boost[tmp]
# print(i, tmp)
print(tmp)
ret[i] = lv.rest_vector(lv.neg(p4_all[tmp]), ret[i])
ret2 = {}
mask = tf.expand_dims(mask, -1)
for i in ret:
ret2[i] = tf.where(mask, ret[i], data["particle"][tp_map[i]]["p"])
print("ret2:", {k: v.shape for k, v in ret2.items()})
# print({i: data["particle"][tp_map[i]]["p"] for i in decay_chain.outs})
pi = np.stack([ret2[i] for i in dat_order], axis=-2)
pi = np.transpose(pi, (1, 0, 2, 3))
total_weights = np.transpose(total_weights.numpy(), (1, 0))
print(pi.shape)
return pi, total_weights
