def test__rambo_generate_all_legs(): delta_cut = 0.01 rambo = Rambo(num_jets=3, num_points=100, w=500., delta_cut=delta_cut, all_legs=True) cut_mom = rambo.generate() assert np.isclose(dot(cut_mom[0][0], cut_mom[0][0]), 0.) assert np.isclose(dot(cut_mom[0][1], cut_mom[0][1]), 0.) assert np.isclose(dot(cut_mom[0][2], cut_mom[0][2]), 0.) assert np.isclose(dot(cut_mom[0][3], cut_mom[0][3]), 0.) assert np.isclose(dot(cut_mom[0][4], cut_mom[0][4]), 0.) for i in cut_mom: too_close, distance = check_all(mom=i, delta=delta_cut, s_com=dot(cut_mom[0][0], cut_mom[0][1]), all_legs=True) assert too_close == False assert distance > delta_cut
def cut_near_split(self, delta_cut, delta_near, return_indices=False): ''' Split momenta into near and cut arrays - near is the region close to the PS cuts and the cut region is the rest of the cut PS :param delta_cut: the PS cut delta :param delta_near: the secondary 'cut' defining the region 'close to' the cut boundary ''' cut_momenta = [] self.near_momenta = [] cut_labels = [] self.near_labels = [] cut_indices = [] near_indices = [] for idx, i in tqdm(enumerate(self.momenta), total=len(self.momenta)): close, min_distance = check_all(mom=i, delta=delta_cut, s_com=dot(i[0], i[1]), all_legs=self.all_legs) if not close: if min_distance < delta_near: self.near_momenta.append(i) self.near_labels.append(self.labels[idx]) near_indices.append(idx) else: cut_momenta.append(i) cut_labels.append(self.labels[idx]) cut_indices.append(idx) if return_indices: return cut_momenta, self.near_momenta, cut_labels, self.near_labels, cut_indices, near_indices else: return cut_momenta, self.near_momenta, cut_labels, self.near_labels
def test__rambo_generate_piecewise(): delta_cut = 0.01 delta_near = 0.02 rambo = Rambo(num_jets=3, num_points=100, w=500., delta_cut=delta_cut, all_legs=False) cut_mom, near_mom = rambo.generate_piecewise(delta_near=delta_near) assert np.isclose(dot(cut_mom[0][0], cut_mom[0][0]), 0.) assert np.isclose(dot(cut_mom[0][1], cut_mom[0][1]), 0.) assert np.isclose(dot(cut_mom[0][2], cut_mom[0][2]), 0.) assert np.isclose(dot(cut_mom[0][3], cut_mom[0][3]), 0.) assert np.isclose(dot(cut_mom[0][4], cut_mom[0][4]), 0.) assert np.isclose(dot(near_mom[0][0], near_mom[0][0]), 0.) assert np.isclose(dot(near_mom[0][1], near_mom[0][1]), 0.) assert np.isclose(dot(near_mom[0][2], near_mom[0][2]), 0.) assert np.isclose(dot(near_mom[0][3], near_mom[0][3]), 0.) assert np.isclose(dot(near_mom[0][4], near_mom[0][4]), 0.) for i in cut_mom: too_close, distance = check_all(mom=i, delta=delta_cut, s_com=dot(cut_mom[0][0], cut_mom[0][1]), all_legs=False) assert too_close == False assert distance > delta_cut assert distance > delta_near for i in near_mom: too_close, distance = check_all(mom=i, delta=delta_cut, s_com=dot(cut_mom[0][0], cut_mom[0][1]), all_legs=False) assert too_close == False assert distance > delta_cut assert distance < delta_near