Skip to content

Commit

Permalink
Strict UI (#17)
Browse files Browse the repository at this point in the history 
Strict UI
  • Loading branch information
@jonas-eschle
jonas-eschle committed May 8, 2019
2 parents 311ce8e + fa4e5aa commit bc32651
Show file tree
Hide file tree
Showing 5 changed files with 109 additions and 79 deletions.
79 changes: 56 additions & 23 deletions phasespace/phasespace.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,11 +87,12 @@ class Particle:
"""

def __init__(self, name, mass=None):
def __init__(self, name, mass): # noqa
self.name = name
self.children = []
if mass is not None and not callable(mass) and not tf.contrib.framework.is_tensor(mass):
mass = tf.constant(mass, dtype=tf.float64)
if not callable(mass) and not isinstance(mass, tf.Variable):
mass = tf.convert_to_tensor(mass, preferred_dtype=tf.float64)
mass = tf.cast(mass, tf.float64)
self._mass = mass

def _do_names_clash(self, particles):
Expand All @@ -114,13 +115,14 @@ def get_mass(self, min_mass=None, max_mass=None, n_events=None):
"""Get the particle mass.
If the particle is resonant, the mass function will be called with the
`min_mass` and `max_mass` parameters.
`min_mass`, `max_mass` and `n_events` parameters.
Arguments:
min_mass (tensor): Lower mass range. Defaults to None, which
is only valid in the case of fixed mass.
max_maxx (tensor): Upper mass range. Defaults to None, which
max_mass (tensor): Upper mass range. Defaults to None, which
is only valid in the case of fixed mass.
n_events (`tf.Tensor`): Number of events to produce. Has to be specified if the particle is resonant.
Return:
tensor: Mass.
Expand Down Expand Up @@ -229,7 +231,8 @@ def _preprocess(momentum, n_events):
momentum = tf.expand_dims(momentum, axis=-1)
return momentum, n_events

def _get_w_max(self, available_mass, masses):
@staticmethod
def _get_w_max(available_mass, masses):
emmax = available_mass + masses[0, :]
emmin = _ZERO
w_max = _ONE
Expand Down Expand Up @@ -358,7 +361,14 @@ def recurse_stable(part):
for part in generated_particles]
return tf.reshape(weights, (n_events,)), w_max, generated_particles, masses

def _recursive_generate(self, momentum, n_events, recalculate_max_weights):
def _recursive_generate(self, n_events, boost_to=None, recalculate_max_weights=False):
if boost_to is not None:
momentum = boost_to
else:
if self.has_fixed_mass:
momentum = tf.stack((0.0, 0.0, 0.0, self.get_mass()), axis=-1)
else:
raise ValueError("Cannot use resonance as top particle")
weights, weights_max, parts, children_masses = self._generate(momentum, n_events)
output_particles = {child.name: parts[child_num]
for child_num, child in enumerate(self.children)}
Expand All @@ -367,7 +377,9 @@ def _recursive_generate(self, momentum, n_events, recalculate_max_weights):
for child_num, child in enumerate(self.children):
if child.has_children:
child_weights, _, child_gen_particles, child_masses = \
child._recursive_generate(parts[child_num], n_events, False)
child._recursive_generate(n_events=n_events,
boost_to=parts[child_num],
recalculate_max_weights=False)
weights *= child_weights
output_particles.update(child_gen_particles)
output_masses.update(child_masses)
Expand Down Expand Up @@ -429,17 +441,22 @@ def recurse_w_max(parent_mass, current_mass_tree):
recurse_w_max(kin.mass(momentum), mass_tree[self.name])
return weights, weights_max, output_particles, output_masses

def generate_unnormalized(self, momentum, n_events=None):
def generate_unnormalized(self, n_events=None, boost_to=None):
"""Generate unnormalized n-body phase space.
Note:
In this method, the event weights and their normalization (the maximum weight)
are returned separately.
Note:
If nor `n_events` nor `boost_to` is given, a single event is generated in the
rest frame of the particle.
Arguments:
`momentum`: Momentum vector of shape (4, x), where x is optional.
`n_events` (optional): Number of events to generate. If `None` (default),
the number of events to generate is calculated from the shape of `momentum`.
the number of events to generate is calculated from the shape of `boost`.
`boost_to`: Momentum vector of shape (4, x), where x is optional, where
the resulting events will be boosted to.
Return:
tuple: Event weights tensor of shape (n_events, ), max event weights tensor, of shape
Expand All @@ -450,24 +467,36 @@ def generate_unnormalized(self, momentum, n_events=None):
tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden.
"""
if not (isinstance(n_events, tf.Variable) or n_events is None):
if n_events is None:
if boost_to is None or boost_to.shape.ndims == 1:
n_events = 1
else:
n_events = tf.shape(boost_to)[1]
if not isinstance(n_events, tf.Variable):
n_events = tf.convert_to_tensor(n_events, preferred_dtype=tf.int32)
n_events = tf.cast(n_events, dtype=tf.int32)

weights, weights_max, parts, _ = self._recursive_generate(momentum=momentum, n_events=n_events,
weights, weights_max, parts, _ = self._recursive_generate(n_events=n_events,
boost_to=boost_to,
recalculate_max_weights=self.has_grandchildren)
return weights, weights_max, parts

def generate(self, momentum, n_events=None):
def generate(self, n_events=None, boost_to=None):
"""Generate normalized n-body phase space.
Events are generated in the rest frame of the particle, unless `boost_to` is given.
Note:
In this method, the event weights are returned normalized to their maximum.
Note:
If nor `n_events` nor `boost_to` is given, a single event is generated in the
rest frame of the particle.
Arguments:
`momentum`: Momentum vector of shape (4, x), where x is optional.
`n_events` (optional): Number of events to generate. If `None` (default),
the number of events to generate is calculated from the shape of `momentum`.
the number of events to generate is calculated from the shape of `boost`.
`boost_to`: Momentum vector of shape (4, x), where x is optional, where
the resulting events will be boosted to.
Return:
tuple: Normalized event weights tensor of shape (n_events, ), and generated
Expand All @@ -478,27 +507,31 @@ def generate(self, momentum, n_events=None):
tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden.
"""
weights, weights_max, parts = self.generate_unnormalized(momentum, n_events)
if n_events is None and boost_to is None:
n_events = 1
weights, weights_max, parts = self.generate_unnormalized(n_events, boost_to)
return weights / weights_max, parts


def generate(p_top, masses, n_events=None):
def generate(mass_top, masses, n_events=None, boost_to=None):
"""Generate an n-body phasespace.
Internally, this function uses `Particle` with a single generation of children.
Arguments:
p_top (Tensor, list): Momentum of the top particle. Can be a list of 4-vectors.
mass_top (`tf.Tensor`, list): Mass of the top particle. Can be a list of 4-vectors.
masses (list): Masses of the child particles.
n_events (int, optional): Number of samples to generate. If n_events is None,
the number of events is deduced from `p_top`.
the number of events is deduced from `mass_top`.
boost_to (`tf.Tensor`): Momentum of the top particle.
Return:
Tensor: 4-momenta of the generated particles, with shape (4xn_particles, n_events).
"""
top = Particle('top').set_children(*[Particle(str(num + 1), mass=mass) for num, mass in enumerate(masses)])
norm_weights, parts = top.generate(p_top, n_events=n_events)
top = Particle('top', mass_top).set_children(*[Particle(str(num + 1), mass=mass)
for num, mass in enumerate(masses)])
norm_weights, parts = top.generate(n_events=n_events, boost_to=boost_to)
return norm_weights, [parts[child.name] for child in top.children]

# EOF
22 changes: 10 additions & 12 deletions tests/helpers/decays.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,9 @@
import tensorflow_probability as tfp

from phasespace import Particle
from phasespace.kinematics import mass

# Use RapidSim values (https://github.com/gcowan/RapidSim/blob/master/config/particles.dat)
B0_MASS = 5279.58
B0_AT_REST = tf.stack((0.0, 0.0, 0.0, B0_MASS), axis=-1)
PION_MASS = 139.57018
KAON_MASS = 493.677
K1_MASS = 1272
Expand All @@ -40,10 +38,10 @@ def kstar_mass(min_mass, max_mass, n_events):
high=max_mass).sample()
return kstar_mass

return Particle('B0').set_children(Particle('K*0', mass=kstar_mass)
.set_children(Particle('K+', mass=KAON_MASS),
Particle('pi-', mass=PION_MASS)),
Particle('gamma', mass=0.0))
return Particle('B0', B0_MASS).set_children(Particle('K*0', mass=kstar_mass)
.set_children(Particle('K+', mass=KAON_MASS),
Particle('pi-', mass=PION_MASS)),
Particle('gamma', mass=0.0))


def bp_to_k1_kstar_pi_gamma(k1_width=K1_WIDTH, kstar_width=KSTARZ_WIDTH):
Expand All @@ -67,11 +65,11 @@ def k1_mass(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 Particle('B+').set_children(Particle('K1+', mass=k1_mass)
.set_children(Particle('K*0', mass=kstar_mass)
.set_children(Particle('K+', mass=KAON_MASS),
Particle('pi-', mass=PION_MASS)),
Particle('pi+', mass=PION_MASS)),
Particle('gamma', mass=0.0))
return Particle('B+', B0_MASS).set_children(Particle('K1+', mass=k1_mass)
.set_children(Particle('K*0', mass=kstar_mass)
.set_children(Particle('K+', mass=KAON_MASS),
Particle('pi-', mass=PION_MASS)),
Particle('pi+', mass=PION_MASS)),
Particle('gamma', mass=0.0))

# EOF
24 changes: 12 additions & 12 deletions tests/test_chain.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,26 +25,26 @@ def test_name_clashes():
"""Test clashes in particle naming."""
# In children
with pytest.raises(KeyError):
Particle('Top').set_children(Particle('Kstarz', mass=decays.KSTARZ_MASS),
Particle('Kstarz', mass=decays.KSTARZ_MASS))
Particle('Top', 0).set_children(Particle('Kstarz', mass=decays.KSTARZ_MASS),
Particle('Kstarz', mass=decays.KSTARZ_MASS))
# With itself
with pytest.raises(KeyError):
Particle('Top').set_children(Particle('Top', mass=decays.KSTARZ_MASS),
Particle('Kstarz', mass=decays.KSTARZ_MASS))
Particle('Top', 0).set_children(Particle('Top', mass=decays.KSTARZ_MASS),
Particle('Kstarz', mass=decays.KSTARZ_MASS))
# In grandchildren
with pytest.raises(KeyError):
Particle('Top').set_children(Particle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(Particle('K+', mass=decays.KAON_MASS),
Particle('pi-', mass=decays.PION_MASS)),
Particle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(Particle('K+', mass=decays.KAON_MASS),
Particle('pi-_1', mass=decays.PION_MASS)))
Particle('Top', 0).set_children(Particle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(Particle('K+', mass=decays.KAON_MASS),
Particle('pi-', mass=decays.PION_MASS)),
Particle('Kstarz0', mass=decays.KSTARZ_MASS)
.set_children(Particle('K+', mass=decays.KAON_MASS),
Particle('pi-_1', mass=decays.PION_MASS)))


def test_kstargamma():
"""Test B0 -> K*gamma."""
with tf.Session() as sess:
norm_weights, particles = sess.run(decays.b0_to_kstar_gamma().generate(decays.B0_AT_REST, 1000))
norm_weights, particles = sess.run(decays.b0_to_kstar_gamma().generate(n_events=1000))
assert len(norm_weights) == 1000
assert all([weight < 1 for weight in norm_weights])
assert len(particles) == 4
Expand All @@ -55,7 +55,7 @@ def test_kstargamma():
def test_k1gamma():
"""Test B+ -> K1 (K*pi) gamma."""
with tf.Session() as sess:
norm_weights, particles = sess.run(decays.bp_to_k1_kstar_pi_gamma().generate(decays.B0_AT_REST, 1000))
norm_weights, particles = sess.run(decays.bp_to_k1_kstar_pi_gamma().generate(n_events=1000))
assert len(norm_weights) == 1000
assert all([weight < 1 for weight in norm_weights])
assert len(particles) == 6
Expand Down
49 changes: 24 additions & 25 deletions tests/test_generate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,15 +20,14 @@

from .helpers import decays

B_AT_REST = decays.B0_AT_REST
BS_AT_REST = tf.stack((0.0, 0.0, 0.0, decays.B0_MASS + 86.8), axis=-1)
B0_MASS = decays.B0_MASS
PION_MASS = decays.PION_MASS


def test_one_event():
"""Test B->pi pi pi."""
with tf.Session() as sess:
norm_weights, particles = sess.run(phasespace.generate(B_AT_REST,
norm_weights, particles = sess.run(phasespace.generate(B0_MASS,
[PION_MASS, PION_MASS, PION_MASS]))
assert len(norm_weights) == 1
assert all([weight < 1 for weight in norm_weights])
Expand All @@ -42,7 +41,7 @@ def test_n_events(n_events):
with tf.Session() as sess:
if isinstance(n_events, tf.Variable):
sess.run(n_events.initializer)
norm_weights, particles = sess.run(phasespace.generate(B_AT_REST,
norm_weights, particles = sess.run(phasespace.generate(B0_MASS,
[PION_MASS, PION_MASS, PION_MASS],
n_events=n_events))
assert len(norm_weights) == 5
Expand All @@ -51,27 +50,27 @@ def test_n_events(n_events):
assert all([part.shape == (4, 5) for part in particles])


def test_n_events_implicit_parent():
"""Test multiple events by passing mutiple parent momenta."""
with tf.Session() as sess:
norm_weights, particles = sess.run(phasespace.generate([B_AT_REST, BS_AT_REST],
[PION_MASS, PION_MASS, PION_MASS]))
assert len(norm_weights) == 2
assert all([weight < 1 for weight in norm_weights])
assert len(particles) == 3
assert all([part.shape == (4, 2) for part in particles])


def test_input_inconsistencies():
"""Test input inconsistencies."""
with tf.Session() as sess:
with pytest.raises(tf.errors.InvalidArgumentError):
sess.run(phasespace.generate([B_AT_REST, BS_AT_REST],
[PION_MASS, PION_MASS, PION_MASS],
n_events=5))
with pytest.raises(tf.errors.InvalidArgumentError):
sess.run(phasespace.generate(B_AT_REST,
[6000.0, PION_MASS, PION_MASS]))
# def test_n_events_implicit_parent():
# """Test multiple events by passing mutiple parent momenta."""
# with tf.Session() as sess:
# norm_weights, particles = sess.run(phasespace.generate([B_AT_REST, BS_AT_REST],
# [PION_MASS, PION_MASS, PION_MASS]))
# assert len(norm_weights) == 2
# assert all([weight < 1 for weight in norm_weights])
# assert len(particles) == 3
# assert all([part.shape == (4, 2) for part in particles])


# def test_input_inconsistencies():
# """Test input inconsistencies."""
# with tf.Session() as sess:
# with pytest.raises(tf.errors.InvalidArgumentError):
# sess.run(phasespace.generate([B_AT_REST, BS_AT_REST],
# [PION_MASS, PION_MASS, PION_MASS],
# n_events=5))
# with pytest.raises(tf.errors.InvalidArgumentError):
# sess.run(phasespace.generate(B_AT_REST,
# [6000.0, PION_MASS, PION_MASS]))


if __name__ == "__main__":
Expand Down
14 changes: 7 additions & 7 deletions tests/test_physics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ def run_test(n_particles, test_prefix):
n_events = tf.Variable(initial_value=first_run_n_events, dtype=tf.int64, use_resource=True)
sess.run(n_events.initializer)

generate = phasespace.generate(decays.B0_AT_REST,
generate = phasespace.generate(decays.B0_MASS,
[decays.PION_MASS] * n_particles,
n_events)
weights1, particles1 = sess.run(generate) # only generate to test change in n_events
Expand Down Expand Up @@ -138,14 +138,14 @@ def test_four_body():
def run_kstargamma(input_file, kstar_width, b_at_rest, suffix):
n_events = 1000000
if b_at_rest:
input_bs = decays.B0_AT_REST
booster = None
rapidsim_getter = rapidsim.get_tree_in_b_rest_frame
else:
input_bs = rapidsim.generate_fonll(decays.B0_MASS, 7, 'b', n_events)
booster = rapidsim.generate_fonll(decays.B0_MASS, 7, 'b', n_events)
rapidsim_getter = rapidsim.get_tree
with tf.Session() as sess:
norm_weights, particles = sess.run(
decays.b0_to_kstar_gamma(kstar_width=kstar_width).generate(input_bs, n_events))
decays.b0_to_kstar_gamma(kstar_width=kstar_width).generate(n_events=n_events, boost_to=booster))
rapidsim_parts = rapidsim_getter(os.path.join(BASE_PATH,
'data',
input_file),
Expand Down Expand Up @@ -211,15 +211,15 @@ def test_kstargamma_resonant_at_rest():
def run_k1_gamma(input_file, k1_width, kstar_width, b_at_rest, suffix):
n_events = 1000000
if b_at_rest:
input_bs = decays.B0_AT_REST
booster = None
rapidsim_getter = rapidsim.get_tree_in_b_rest_frame
else:
input_bs = rapidsim.generate_fonll(decays.B0_MASS, 7, 'b', n_events)
booster = rapidsim.generate_fonll(decays.B0_MASS, 7, 'b', n_events)
rapidsim_getter = rapidsim.get_tree
with tf.Session() as sess:
norm_weights, particles = sess.run(
decays.bp_to_k1_kstar_pi_gamma(k1_width=k1_width, kstar_width=kstar_width)
.generate(input_bs, n_events))
.generate(n_events=n_events, boost_to=booster))
rapidsim_parts = rapidsim_getter(
os.path.join(BASE_PATH,
'data',
Expand Down

0 comments on commit bc32651

Please sign in to comment.