Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
zero_component,
tf.sqrt(tf.square(pds[part_num - 1]) +
tf.square(tf.gather(masses, [part_num], axis=1)))],
axis=1))
with tf.control_dependencies([n_events]):
cos_z = (tf.constant(2.0, dtype=tf.float64) * tf.random.uniform((n_events, 1), dtype=tf.float64)
- tf.constant(1.0, dtype=tf.float64))
sin_z = tf.sqrt(tf.constant(1.0, dtype=tf.float64) - cos_z * cos_z)
ang_y = (tf.constant(2.0, dtype=tf.float64) * tf.constant(pi, dtype=tf.float64)
* tf.random.uniform((n_events, 1), dtype=tf.float64))
cos_y = tf.math.cos(ang_y)
sin_y = tf.math.sin(ang_y)
# Do the rotations
for j in range(part_num + 1):
px = kin.x_component(generated_particles[j])
py = kin.y_component(generated_particles[j])
# Rotate about z
# TODO(Mayou36): only list? will be overwritten below anyway, but can `*_component` handle it?
generated_particles[j] = tf.concat([cos_z * px - sin_z * py,
sin_z * px + cos_z * py,
kin.z_component(generated_particles[j]),
kin.time_component(generated_particles[j])],
axis=1)
# Rotate about y
px = kin.x_component(generated_particles[j])
pz = kin.z_component(generated_particles[j])
generated_particles[j] = tf.concat([cos_y * px - sin_y * pz,
kin.y_component(generated_particles[j]),
sin_y * px + cos_y * pz,
kin.time_component(generated_particles[j])],
axis=1)
if part_num == (n_particles - 1):
# Do the rotations
for j in range(part_num + 1):
px = kin.x_component(generated_particles[j])
py = kin.y_component(generated_particles[j])
# Rotate about z
# TODO(Mayou36): only list? will be overwritten below anyway, but can `*_component` handle it?
generated_particles[j] = tf.concat([cos_z * px - sin_z * py,
sin_z * px + cos_z * py,
kin.z_component(generated_particles[j]),
kin.time_component(generated_particles[j])],
axis=1)
# Rotate about y
px = kin.x_component(generated_particles[j])
pz = kin.z_component(generated_particles[j])
generated_particles[j] = tf.concat([cos_y * px - sin_y * pz,
kin.y_component(generated_particles[j]),
sin_y * px + cos_y * pz,
kin.time_component(generated_particles[j])],
axis=1)
if part_num == (n_particles - 1):
break
betas = (pds[part_num] / tf.sqrt(tf.square(pds[part_num]) + tf.square(inv_masses[part_num])))
generated_particles = [kin.lorentz_boost(part,
tf.concat([zero_component,
betas,
zero_component],
axis=1))
for part in generated_particles]
part_num += 1
# Final boost of all particles
generated_particles = [kin.lorentz_boost(part, p_top_boost)
for part in generated_particles]