def test_custom_weights(self): lbn = LBN(10, boost_mode=LBN.PAIRS, particle_weights=self.custom_particle_weights, restframe_weights=self.custom_restframe_weights) lbn(self.vectors_t, features=self.feature_set).numpy() self.assertEqual(lbn.particle_weights.numpy().shape, (10, 10)) self.assertEqual(lbn.restframe_weights.numpy().shape, (10, 10)) self.assertEqual(np.sum(lbn.particle_weights.numpy()), 3) self.assertEqual(np.sum(lbn.restframe_weights.numpy()), 3) # compare sum of vector components of first combined particles and restframes in batch pos 1 target_particle_sum = np.sum(self.vectors[1, 0] + self.vectors[1, 1]) target_restframe_sum = np.sum(self.vectors[1, 1] + self.vectors[1, 2]) self.assertAlmostEqual(np.sum(lbn.particles.numpy()[1, 0]), target_particle_sum, 3) self.assertAlmostEqual(np.sum(lbn.restframes.numpy()[1, 0]), target_restframe_sum, 3) # test wrong shape lbn = LBN(10, boost_mode=LBN.PAIRS, particle_weights=self.custom_particle_weights, restframe_weights=self.custom_restframe_weights[:-1]) with self.assertRaises(ValueError): lbn(self.vectors_t, features=self.feature_set).numpy()
def test_custom_feature_factory(self): class MyFeatureFactory(FeatureFactory): def px_plus_py(self): return self.px() + self.py() lbn = LBN(10, boost_mode=LBN.PAIRS, feature_factory=MyFeatureFactory) self.assertIn("px_plus_py", lbn.available_features) with self.assertRaises(TypeError): LBN(10, boost_mode=LBN.PAIRS, feature_factory="foo")
def test_boosting_combinations(self): lbn = LBN(10, boost_mode=LBN.COMBINATIONS, particle_weights=self.custom_particle_weights) lbn(self.vectors_t, features=self.feature_set).numpy() # compare all components of the first boosted particle in batch pos 1 # see test_boosting_pairs for manual boost computation p1 = lbn.particles.numpy()[1, 0] components = list(self.vectors[1, 0] + self.vectors[1, 1]) for i, v in enumerate(components): self.assertAlmostEqual(p1[i], v, 3) p2 = lbn.particles.numpy()[1, 1] components = list(self.vectors[1, 0]) for i, v in enumerate(components): self.assertAlmostEqual(p2[i], v, 5) # boosted particle 0 is p1 boosted into p2 boosted = lbn.boosted_particles.numpy()[1, 0] components = [288.7326, 172.70781, 102.427, 146.44083] for i, v in enumerate(components): self.assertAlmostEqual(boosted[i], v, 3) # boosted particle 45 is p2 boosted into p1 boosted = lbn.boosted_particles.numpy()[1, 45] components = [69.299545, -19.58605, -18.497059, -53.21913] for i, v in enumerate(components): self.assertAlmostEqual(boosted[i], v, 3)
def test_boosting_pairs(self): lbn = LBN(10, boost_mode=LBN.PAIRS, particle_weights=self.custom_particle_weights, restframe_weights=self.custom_restframe_weights) lbn(self.vectors_t, features=self.feature_set) # compare all components of the first boosted particle in batch pos 1 particle = lbn.particles.numpy()[1, 0] components = list(self.vectors[1, 0] + self.vectors[1, 1]) for i, v in enumerate(components): self.assertAlmostEqual(particle[i], v, 3) restframe = lbn.restframes.numpy()[1, 0] components = list(self.vectors[1, 1] + self.vectors[1, 2]) for i, v in enumerate(components): self.assertAlmostEqual(restframe[i], v, 3) # boosted values computed ROOT TLorentzVector's via # p = TLorentzVector(particle[1], particle[2], particle[3], particle[0]) # r = TLorentzVector(restframe[1], restframe[2], restframe[3], restframe[0]) # p = p.Boost(-r.BoostVector()) boosted = lbn.boosted_particles.numpy()[1, 0] components = [217.82007, -93.470245, 56.69007, -117.862404] for i, v in enumerate(components): self.assertAlmostEqual(boosted[i], v, 4)
def _test_features(self, autograph): lbn = LBN(10, boost_mode=LBN.PAIRS, particle_weights=self.custom_particle_weights, restframe_weights=self.custom_restframe_weights) # add a custom feature @lbn.register_feature def px_plus_py(factory): return factory.px() + factory.py() all_features = [ "E", "px", "py", "pz", "pt", "p", "m", "phi", "eta", "beta", "gamma", "pair_cos", "pair_dr", "pair_ds", "pair_dy", "px_plus_py", ] self.assertEqual(set(lbn.available_features), set(all_features)) features = lbn(self.vectors_t, features=all_features, autograph=autograph).numpy() # perform tests at batch pos 1 features = features[1] # make all tests on the first boosted particle at batch pos 1 self.assertAlmostEqual(features[0], 217.82007, 4) self.assertAlmostEqual(features[10], -93.470245, 4) self.assertAlmostEqual(features[20], 56.69007, 4) self.assertAlmostEqual(features[30], -117.862404, 4) self.assertAlmostEqual(features[40], 109.318115, 4) self.assertAlmostEqual(features[50], 160.75446, 4) self.assertAlmostEqual(features[60], 146.98158, 4) self.assertAlmostEqual(features[70], 2.5964046, 4) self.assertAlmostEqual(features[80], -0.9355755, 4) self.assertAlmostEqual(features[90], 0.7380149, 4) self.assertAlmostEqual(features[100], 1.4819548, 4) # test pairwise features w.r.t. boosted particle 2, i.e., feature pos 0 self.assertAlmostEqual(features[110], 0.64787644, 4) self.assertAlmostEqual(features[155], 2.6730149, 4) self.assertAlmostEqual(features[200], -136.8383, 4) self.assertAlmostEqual(features[245], -1.3652772, 4) # test the custom feature self.assertAlmostEqual(features[290], -36.780174, 4)
def test_register_feature(self): lbn = LBN(10, boost_mode=LBN.PAIRS) self.assertNotIn("px_plus_py", lbn.available_features) @lbn.register_feature def px_plus_py(factory): return factory.px() + factory.py() self.assertIn("px_plus_py", lbn.available_features)
def test_external_features(self): lbn = LBN(10, boost_mode=LBN.PAIRS) ext = tf.Variable([[1, 2], [3, 4]], dtype=tf.float32) features = lbn(self.vectors_t, features=self.feature_set, external_features=ext).numpy() self.assertEqual(features.shape[1], lbn.n_features) self.assertEqual(features.shape, (2, 97))
def test_aux_features(self): lbn = LBN(10, boost_mode=LBN.PAIRS) features = lbn(self.vectors_aux_t, features=self.feature_set).numpy() self.assertEqual(lbn.n_dim, 6) self.assertEqual(lbn.n_aux, 2) self.assertEqual(lbn.n_auxiliaries, 10) self.assertEqual(lbn.aux_weights.shape, (10, 10, 2)) self.assertEqual(features.shape[1], lbn.n_features) self.assertEqual(features.shape, (2, 115))
def test_constructor_boost_mode_pairs(self): lbn = LBN(10, boost_mode=LBN.PAIRS) self.assertEqual(lbn.n_particles, 10) self.assertEqual(lbn.n_restframes, 10) self.assertEqual(lbn.n_out, 10) self.assertIsNone(lbn.n_features) features = lbn(self.vectors_t, features=self.feature_set).numpy() self.assertEqual(lbn.n_in, 10) self.assertEqual(features.shape[1], lbn.n_features) self.assertEqual(features.shape, (2, 95))
def test_constructor_boost_mode_product(self): lbn = LBN(10, 4, boost_mode=LBN.PRODUCT) self.assertEqual(lbn.n_particles, 10) self.assertEqual(lbn.n_restframes, 4) self.assertEqual(lbn.n_out, 40) self.assertIsNone(lbn.n_features) features = lbn(self.vectors_t, features=self.feature_set).numpy() self.assertEqual(lbn.n_in, 10) self.assertEqual(features.shape[1], lbn.n_features) self.assertEqual(features.shape, (2, 980))
def test_constructor_boost_mode_combinations(self): lbn = LBN(10, boost_mode=LBN.COMBINATIONS) self.assertEqual(lbn.n_particles, 10) self.assertEqual(lbn.n_restframes, 10) self.assertEqual(lbn.n_out, 90) self.assertIsNone(lbn.n_features) features = lbn(self.vectors_t, features=self.feature_set).numpy() self.assertEqual(lbn.n_in, 10) self.assertEqual(features.shape[1], lbn.n_features) self.assertEqual(features.shape, (2, 4455))
def test_boosting_product(self): lbn = LBN(10, 4, boost_mode=LBN.PRODUCT, particle_weights=self.custom_particle_weights, restframe_weights=self.custom_restframe_weights[:, :4]) lbn(self.vectors_t, features=self.feature_set).numpy() # compare all components of the first boosted particle in batch pos 1 # see test_boosting_pairs for manual boost computation boosted = lbn.boosted_particles.numpy()[1, 0] components = [217.82007, -93.470245, 56.69007, -117.862404] for i, v in enumerate(components): self.assertAlmostEqual(boosted[i], v, 4)
def test_post_build_attributes(self): attrs = [ "particle_weights", "abs_particle_weights", "clip_particle_weights", "restframe_weights", "abs_restframe_weights", "clip_restframe_weights", "aux_weights", "n_in", "n_dim", "n_aux", "I", "U", "inputs", "inputs_E", "inputs_px", "inputs_py", "inputs_pz", "particles_E", "particles_px", "particles_py", "particles_pz", "inputs_aux", "particles_pvec", "particles", "restframes_E", "restframes_px", "restframes_py", "restframes_pz", "restframes_pvec", "restframes", "Lambda", "boosted_particles", "boosted_features", "aux_features", "features", ] lbn = LBN(10, boost_mode=LBN.PAIRS) for attr in attrs: self.assertIn(getattr(lbn, attr), (None, True, False)) lbn(self.vectors_aux_t, features=self.feature_set).numpy() for attr in attrs: self.assertIsNotNone(getattr(lbn, attr), None)
def test_feature_caching(self): class MyFeatureFactory(FeatureFactory): def __init__(self, *args, **kwargs): super(MyFeatureFactory, self).__init__(*args, **kwargs) self.count = 0 def px_plus_py(self): self.count += 1 return self.px() + self.py() lbn = LBN(10, boost_mode=LBN.PAIRS, feature_factory=MyFeatureFactory) self.assertEqual(lbn.feature_factory.count, 0) lbn(self.vectors_t, features=self.feature_set + ["px_plus_py"]).numpy() self.assertEqual(lbn.feature_factory.count, 1) lbn.feature_factory.px_plus_py() self.assertEqual(lbn.feature_factory.count, 1)
def test_feature_caching(self): class MyFeatureFactory(FeatureFactory): def __init__(self, *args, **kwargs): super(MyFeatureFactory, self).__init__(*args, **kwargs) self.count = 0 @FeatureFactory.hidden_feature def _px_plus_py(self, **opts): self.count += 1 return self.px(**opts) + self.py(**opts) @FeatureFactory.single_feature def px_plus_py(self, **opts): return self._px_plus_py(**opts) @FeatureFactory.single_feature def px_plus_py_2(self, **opts): return self._px_plus_py(**opts)**2. lbn = LBN(10, boost_mode=LBN.PAIRS, feature_factory=MyFeatureFactory) self.assertEqual(lbn.feature_factory.count, 0) lbn(self.vectors_t, features=self.feature_set + ["px_plus_py", "px_plus_py_2"]).numpy() self.assertEqual(lbn.feature_factory.count, 1) lbn(self.vectors_t, features=self.feature_set + ["px_plus_py", "px_plus_py_2"]).numpy() self.assertEqual(lbn.feature_factory.count, 2) lbn.feature_factory.px_plus_py(_symbolic=False) self.assertEqual(lbn.feature_factory.count, 2) lbn.feature_factory.clear_eager_cache() lbn.feature_factory.px_plus_py(_symbolic=False) self.assertEqual(lbn.feature_factory.count, 3)
def test_unknown_boost_mode(self): with self.assertRaises(ValueError): LBN(10, boost_mode="definitely_not_there")
# Define placeholders is_training_t = tf.placeholder(tf.bool, name="my_bool") x_t = tf.placeholder(tf.float32, shape=[None, 3, 4], name="x_t") y_t = tf.placeholder(tf.float32, shape=[None], name="y_t") # create `tf.data.Dataset`, build batches and make it iterable dataset = tf.data.Dataset.from_tensor_slices((x_t, y_t)) dataset = dataset.batch(BATCH_SIZE, drop_remainder=True) iterator = tf.data.Iterator.from_structure(dataset.output_types, dataset.output_shapes) training_init = iterator.make_initializer(dataset, name='dataset_init') next_batch = iterator.get_next() # initialize the LBN, set 4 combinations lbn = LBN(6, boost_mode=LBN.PAIRS, batch_norm=True, is_training=is_training_t) # # create a feature tensor based on input four-vectors # # and register it to the feature factory, here: # # m**2 = E**2 - p**2 of all particle combinations # @lbn.register_feature # def pair_m2(factory): # all_pair_m2 = factory.E**2 - factory.p**2 # return tf.gather(tf.reshape(all_pair_m2, [-1, factory.n**2]), factory.triu_indices, axis=1) # extract features out of the lbn features = lbn( next_batch[0], features=["E", "px", "py", "pz", "pt", "eta", "phi", "m", "pair_cos"]) # Performs all the Lorenz Boost Network
def test_pre_build_attributes(self): lbn = LBN(10, boost_mode=LBN.PAIRS) attrs = ["epsilon", "name"] for attr in attrs: self.assertIsNotNone(getattr(lbn, attr))
def test_constructor(self): lbn = LBN(10) self.assertIsInstance(lbn, LBN)