Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
min_mass = tf.cast(min_mass, tf.float64)
max_mass = tf.cast(max_mass, tf.float64)
kstar_width_cast = tf.cast(kstar_width, tf.float64)
kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64)
kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,))
if kstar_width > 0:
kstar_mass = tfp.distributions.TruncatedNormal(loc=kstar_mass,
scale=kstar_width_cast,
low=min_mass,
high=max_mass).sample()
return kstar_mass
return GenParticle('B0', B0_MASS).set_children(GenParticle('K*0', mass=kstar_mass)
.set_children(GenParticle('K+', mass=KAON_MASS),
GenParticle('pi-', mass=PION_MASS)),
GenParticle('gamma', mass=0.0))
def test_name_clashes():
"""Test clashes in particle naming."""
# In children
with pytest.raises(KeyError):
GenParticle('Top', 0).set_children(GenParticle('Kstarz', mass=decays.KSTARZ_MASS),
GenParticle('Kstarz', mass=decays.KSTARZ_MASS))
# With itself
with pytest.raises(KeyError):
GenParticle('Top', 0).set_children(GenParticle('Top', mass=decays.KSTARZ_MASS),
GenParticle('Kstarz', mass=decays.KSTARZ_MASS))
# In grandchildren
with pytest.raises(KeyError):
GenParticle('Top', 0).set_children(GenParticle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(GenParticle('K+', mass=decays.KAON_MASS),
GenParticle('pi-', mass=decays.PION_MASS)),
GenParticle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(GenParticle('K+', mass=decays.KAON_MASS),
GenParticle('pi-_1', mass=decays.PION_MASS)))
masses = tf.broadcast_to(mass, shape=(n_events,))
if kstar_width > 0:
masses = tfp.distributions.TruncatedNormal(loc=masses,
scale=width,
low=min_mass,
high=max_mass).sample()
return masses
def k1_mass(min_mass, max_mass, n_events):
return res_mass(K1_MASS, k1_width, min_mass, max_mass, n_events)
def kstar_mass(min_mass, max_mass, n_events):
return res_mass(KSTARZ_MASS, kstar_width, min_mass, max_mass, n_events)
return GenParticle('B+', B0_MASS).set_children(GenParticle('K1+', mass=k1_mass)
.set_children(GenParticle('K*0', mass=kstar_mass)
.set_children(GenParticle('K+', mass=KAON_MASS),
GenParticle('pi-', mass=PION_MASS)),
GenParticle('pi+', mass=PION_MASS)),
GenParticle('gamma', mass=0.0))
def test_name_clashes():
"""Test clashes in particle naming."""
# In children
with pytest.raises(KeyError):
GenParticle('Top', 0).set_children(GenParticle('Kstarz', mass=decays.KSTARZ_MASS),
GenParticle('Kstarz', mass=decays.KSTARZ_MASS))
# With itself
with pytest.raises(KeyError):
GenParticle('Top', 0).set_children(GenParticle('Top', mass=decays.KSTARZ_MASS),
GenParticle('Kstarz', mass=decays.KSTARZ_MASS))
# In grandchildren
with pytest.raises(KeyError):
GenParticle('Top', 0).set_children(GenParticle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(GenParticle('K+', mass=decays.KAON_MASS),
GenParticle('pi-', mass=decays.PION_MASS)),
GenParticle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(GenParticle('K+', mass=decays.KAON_MASS),
GenParticle('pi-_1', mass=decays.PION_MASS)))
def test_reset_children():
"""Test when children are set twice."""
top = GenParticle('Top', 0).set_children(GenParticle('Child1', mass=decays.KSTARZ_MASS),
GenParticle('Child2', mass=decays.KSTARZ_MASS))
with pytest.raises(ValueError):
top.set_children(GenParticle('Child3', mass=decays.KSTARZ_MASS),
GenParticle('Child4', mass=decays.KSTARZ_MASS))
def test_no_children():
"""Test when no children have been configured."""
top = GenParticle('Top', 0)
with pytest.raises(ValueError):
top.generate(n_events=1)
def test_name_clashes():
"""Test clashes in particle naming."""
# In children
with pytest.raises(KeyError):
GenParticle('Top', 0).set_children(GenParticle('Kstarz', mass=decays.KSTARZ_MASS),
GenParticle('Kstarz', mass=decays.KSTARZ_MASS))
# With itself
with pytest.raises(KeyError):
GenParticle('Top', 0).set_children(GenParticle('Top', mass=decays.KSTARZ_MASS),
GenParticle('Kstarz', mass=decays.KSTARZ_MASS))
# In grandchildren
with pytest.raises(KeyError):
GenParticle('Top', 0).set_children(GenParticle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(GenParticle('K+', mass=decays.KAON_MASS),
GenParticle('pi-', mass=decays.PION_MASS)),
GenParticle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(GenParticle('K+', mass=decays.KAON_MASS),
GenParticle('pi-_1', mass=decays.PION_MASS)))
def test_grandchildren():
"""Test that grandchildren detection is correct."""
top = GenParticle('Top', 0)
assert not top.has_children
assert not top.has_grandchildren
assert not top.set_children(GenParticle('Child1', mass=decays.KSTARZ_MASS),
GenParticle('Child2', mass=decays.KSTARZ_MASS)).has_grandchildren