How to use the phasespace.kinematics function in phasespace

in current_mass_tree.items()
                                                                         if ch_it != child})),
                                               child_mass)
                        masses.append(sum(get_flattened_values(child_mass)))
                    else:
                        masses.append(child_mass)
                masses = tf.concat(masses, axis=1)
                w_max *= self._get_w_max(available_mass, masses)
                return w_max

            mass_tree = {}
            build_mass_tree(self, mass_tree)
            momentum = process_list_to_tensor(momentum)
            if len(momentum.shape.as_list()) == 1:
                momentum = tf.expand_dims(momentum, axis=-1)
            weights_max = tf.reshape(recurse_w_max(kin.mass(momentum), mass_tree[self.name]),
                                     (n_events,))
        return weights, weights_max, output_particles, output_masses

This method generates the same results as the GENBOD routine.

        Arguments:
            momentum (tensor): Momentum of the parent particle. All generated particles
                will be boosted to that momentum.
            n_events (int): Number of events to generate.

        Return:
            tuple: Result of the generation (per-event weights, maximum weights, output particles
                and their output masses).

        """
        if not self.children:
            raise ValueError("No children have been configured")
        p_top, n_events = self._preprocess(momentum, n_events)
        top_mass = tf.broadcast_to(kin.mass(p_top), (n_events, 1))
        n_particles = len(self.children)

        # Prepare masses
        def recurse_stable(part):
            output_mass = 0
            for child in part.children:
                if child.has_fixed_mass:
                    output_mass += child.get_mass()
                else:
                    output_mass += recurse_stable(child)
            return output_mass

        mass_from_stable = tf.broadcast_to(
            tf.reduce_sum([child.get_mass() for child in self.children
                           if child.has_fixed_mass],
                          axis=0),

-pds[part_num - 1],
                                                  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)