How to use phasespace - 10 common examples

def run_test(n_particles, test_prefix):
    sess = tf.Session()
    first_run_n_events = 100
    main_run_n_events = 100000
    n_events = tf.Variable(initial_value=first_run_n_events, dtype=tf.int64, use_resource=True)
    sess.run(n_events.initializer)

    generate = phasespace.nbody_decay(decays.B0_MASS,
                                      [decays.PION_MASS] * n_particles) \
        .generate_tensor(n_events)
    weights1, _ = sess.run(generate)  # only generate to test change in n_events
    assert len(weights1) == first_run_n_events

    # change n_events and run again
    n_events.load(main_run_n_events, session=sess)
    weights, particles = sess.run(generate)
    parts = np.concatenate([particles[f"p_{part_num}"] for part_num in range(n_particles)], axis=1)
    histos = [make_norm_histo(parts[:, coord],
                              range_=(-3000 if coord % 4 != 3 else 0, 3000),
                              weights=weights)
              for coord in range(parts.shape[1])]
    weight_histos = make_norm_histo(weights, range_=(0, 1+1e-8))
    ref_histos, ref_weights = create_ref_histos(n_particles)
    p_values = np.array([ks_2samp(histos[coord], ref_histos[coord])[1]

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

def test_n_events(n_events):
    """Test 5 B->pi pi pi."""
    decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS])
    norm_weights, particles = decay.generate(n_events=n_events)
    assert len(norm_weights) == 5
    assert all([weight < 1 for weight in norm_weights])
    assert len(particles) == 3
    assert all([part.shape == (5, 4) for part in particles.values()])