Aufgrund einer Konfigurationsänderung wird die GitLab Registry ab 10 Uhr nur Read Only zur Verfügung stehen. / Due to a configuration change, the GitLab Registry will be available for read-only access from 10am.

Commit 164bf17b authored by Marcel Rieger's avatar Marcel Rieger

Add initial files.

parent 9b1eaf62
[flake8]
max-line-length = 101
ignore = E306, E402, E128, E722, E731, E741, F812, F822, F403, F401, W504, W605
Copyright (c) 2018-2019, Marcel Rieger
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
* Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
* Neither the name of the copyright holder nor the names of its contributors
may be used to endorse or promote products derived from this software without
specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# Lorentz Boost Network (LBN) [![pipeline status](https://git.rwth-aachen.de/3pia/lbn/badges/master/pipeline.svg)](https://git.rwth-aachen.de/3pia/lbn/commits/master)
TensorFlow implementation of the Lorentz Boost Network from [arXiv:XXXX.XXXXX [hep-ex]](https://arxiv.org/abs/XXXX.XXXXX).
Original repository: [git.rwth-aachen.de/3pia/lbn](https://git.rwth-aachen.de/3pia/lbn)
### Usage example
```python
from lbn import LBN
# initialize the LBN, set 10 combinations and pairwise boosting
lbn = LBN(10, boost_mode=LBN.PAIRS)
# create a feature tensor based on input four-vectors
features = lbn(four_vectors)
# use the features as input for a subsequent, application-specific network
...
```
### Installation and dependencies
Via [pip](https://pypi.python.org/pypi/lbn):
```bash
pip install lbn
```
NumPy and TensorFlow are the only dependencies.
### Testing
Tests should be run for Python 2 and 3. The following commands assume you are root directory of the LBN respository:
```bash
python -m unittest test
# or via docker, python 2
docker run --rm -v `pwd`:/root/lbn -w /root/lbn tensorflow/tensorflow:latest python -m unittest test
# or via docker, python 3
docker run --rm -v `pwd`:/root/lbn -w /root/lbn tensorflow/tensorflow:latest-py3 python -m unittest test
```
### Contributing
If you like to contribute, we are happy to receive pull requests. Just make sure to add new test cases and run the tests. Also, please use a coding style that is compatible with our `.flake8` config.
### Development
- Original source hosted at [RWTH GitLab](https://git.rwth-aachen.de/3pia/lbn)
- Report issues, questions, feature requests on [RWTH GitLab Issues](https://git.rwth-aachen.de/3pia/lbn/issues)
# coding: utf-8
"""
Lorentz Boost Network (LBN) implementation using TensorFlow.
For more information, see https://arxiv.org/abs/XXXX.XXXXX.
"""
__author__ = "Marcel Rieger"
__copyright__ = "Copyright 2018-2019, Marcel Rieger"
__license__ = "BSD"
__credits__ = ["Martin Erdmann", "Erik Geiser", "Yannik Rath", "Marcel Rieger"]
__contact__ = "https://git.rwth-aachen.de/3pia/lbn"
__email__ = "marcel.rieger@cern.ch"
__version__ = "0.0.1"
__all__ = ["LBN", "FeatureFactoryBase", "FeatureFactory"]
import functools
import numpy as np
import tensorflow as tf
class LBN(object):
"""
Lorentz Boost Network (LBN) class.
Usage example:
.. code-block:: python
# initialize the LBN
lbn = LBN(10, boost_mode=LBN.PAIRS)
# create a feature tensor based on input four-vectors
features = lbn(four_vectors)
# use the features as input for a subsequent, application-specific network
...
*n_particles* and *n_restframes* are the number of particle and rest-frame combinations to
build. Their interpretation depends on the *boost_mode*. *n_restframes* is only used for the
*PRODUCT* mode. It is inferred from *n_particles* for *PAIRS* and *COMBINATIONS*.
*is_training* can an external scalar boolean tensor denoting the training or testing phase. If
*None*, in internal flag is created instead and available under the same name. *epsilon* is
supposed to be a small number that is used in various places for numerical stability. *name* is
the main namespace of the LBN and defaults to the class name.
*feature_factory* must be a subclass of :py:class:`FeatureFactoryBase` and provides the
available, generic mappings from boosted particles to output features of the LBN. If *None*, the
default :py:class:`FeatureFactory` is used.
*particle_weights* and *restframe_weights* can refer to externally defined variables with custom
initialized weights. If set, their shape must match the number of combinations and inputs. For
simple initialization tests, *weight_init* can be a tuple containing the Gaussian mean and
standard deviation that is passed to ``tf.random_normal``. When *None*, and the weight tensors
are created internally, mean and standard deviation default to *0* and *1 / combinations*. When
*abs_particle_weights* (*abs_restframe_weights*) is *True*, ``tf.abs`` is applied to the
particle (rest frame) weights. When *clip_particle_weights* (*clip_restframe_weights*) is
*True*, particle (rest frame) weights are clipped at *epsilon*, or at the passed value if it is
not a boolean. Note that the abs operation is applied before clipping.
*batch_norm* defines whether feature scaling via batch normalization with floating averages is
applied to the output features. It can be a single value or a tuple of two values that are
passed to ``tf.layers.batch_normalization``. Make sure to set the proper value for the
*is_training* flag in the feed dict.
Instances of this class store most of the intermediate tensors (such as inputs, combinations
weights, boosted particles, boost matrices, raw features, etc) for later inspection. Note that
most of these tensors are set after :py:meth:`build` (or the :py:meth:`__call__` shorthand as
shown above) are invoked.
"""
# available boost modes
PAIRS = "pairs"
PRODUCT = "product"
COMBINATIONS = "combinations"
def __init__(self, n_particles, n_restframes=None, boost_mode=PAIRS, feature_factory=None,
particle_weights=None, abs_particle_weights=True, clip_particle_weights=False,
restframe_weights=None, abs_restframe_weights=True, clip_restframe_weights=False,
weight_init=None, batch_norm=True, is_training=None, epsilon=1e-5, name=None):
super(LBN, self).__init__()
# determine the number of output particles, which depends on the boost mode
# PAIRS:
# n_restframes set to n_particles, boost pairwise, n_out = n_particles
# PRODUCT:
# boost n_particles into n_restframes, n_out = n_partiles * n_restframes
# COMBINATIONS:
# build only particles, boost them into each other, except for boosts of particles into
# themselves, n_out = n**2 - n
if boost_mode == self.PAIRS:
n_restframes = n_particles
self.n_out = n_particles
elif boost_mode == self.PRODUCT:
self.n_out = n_particles * n_restframes
elif boost_mode == self.COMBINATIONS:
n_restframes = n_particles
self.n_out = n_particles**2 - n_particles
else:
raise ValueError("unknown boost_mode '{}'".format(boost_mode))
# store boost mode and number of particles and restframes to build
self.boost_mode = boost_mode
self.n_particles = n_particles
self.n_restframes = n_restframes
# output batch normalization
if isinstance(batch_norm, bool):
self.batch_norm_center = batch_norm
self.batch_norm_scale = batch_norm
elif isinstance(batch_norm, (list, tuple)) and len(batch_norm) == 2:
self.batch_norm_center, self.batch_norm_scale = batch_norm
else:
raise ValueError("invalid batch_norm, should be bool and list/tuple of two bools")
# particle weights and settings
self.particle_weights = particle_weights
self.abs_particle_weights = abs_particle_weights
self.clip_particle_weights = clip_particle_weights
# rest frame weigths and settings
self.restframe_weights = restframe_weights
self.abs_restframe_weights = abs_restframe_weights
self.clip_restframe_weights = clip_restframe_weights
# custom weight init parameters in a tuple (mean, stddev)
self.weight_init = weight_init
# training flag
self.is_training = is_training
# epsilon for numerical stability
self.epsilon = epsilon
# internal name
self.name = name or self.__class__.__name__
# sizes that are set during build
self.n_in = None # number of input particles
self.n_dim = None # size per input vector, must be four
# constants
self.I = None # the I matrix
self.U = None # the U matrix
# tensor of input vectors
self.inputs = None
# split input tensors
self.inputs_E = None # energy column of inputs
self.inputs_px = None # px column of inputs
self.inputs_py = None # py column of inputs
self.inputs_pz = None # pz column of inputs
# tensors of particle combinations
self.particles_E = None # energy column of combined particles
self.particles_px = None # px column of combined particles
self.particles_py = None # py column of combined particles
self.particles_pz = None # pz column of combined particles
self.particles_pvec = None # p vectors of combined particles
self.particles = None # stacked 4-vectors of combined particles
# tensors of rest frame combinations
self.restframes_E = None # energy column of combined restframes
self.restframes_px = None # px column of combined restframes
self.restframes_py = None # py column of combined restframes
self.restframes_pz = None # pz column of combined restframes
self.restframes_pvec = None # p vectors of combined restframes
self.restframes = None # stacked 4-vectors of combined restframes
# Lorentz boost matrix with shape (batch, n_out, 4, 4)
self.Lambda = None
# boosted particles with shape (batch, n_out, 4)
self.boosted_particles = None
# intermediate features
self._raw_features = None # raw features before batch normalization, etc
self._norm_features = None # features after batch normalization, if used
# final output features
self.features = None
# initialize the feature factory
if feature_factory is None:
feature_factory = FeatureFactory
elif not issubclass(feature_factory, FeatureFactoryBase):
raise TypeError("feature_factory '{}' is not a subclass of FeatureFactoryBase".format(
feature_factory))
self.feature_factory = feature_factory(self)
@property
def available_features(self):
"""
Shorthand to access the list of available features in the :py:attr:`feature_factory`.
"""
return list(self.feature_factory._feature_funcs.keys())
@property
def n_features(self):
"""
Returns the number of created output features which depends on the number of boosted
particles and the feature set.
"""
if self.features is None:
return None
return self.features.shape[-1].value
def register_feature(self, func=None, **kwargs):
"""
Shorthand to register a new feautre to the current :py:attr:`feature_factory` instance. Can
be used as a (configurable) decorator. The decorated function receives the feature factory
instance as the only argument. All *kwargs* are forwarded to
:py:meth:`FeatureFactoryBase._wrap_feature`. Example:
.. code-block:: python
lbn = LBN(10, boost_mode=LBN.PAIRS)
@lbn.register_feature
def px_plus_py(ff):
return ff.px() + ff.py()
print("px_plus_py" in lbn.available_features) # -> True
# or register with a different name
@lbn.register_feature(name="pxy")
def px_plus_py(ff):
return ff.px() + ff.py()
print("pxy" in lbn.available_features) # -> True
"""
def decorator(func):
return self.feature_factory._wrap_feature(func, **kwargs)
return decorator(func) if func else decorator
def __call__(self, *args, **kwargs):
"""
Shorthand for :py:meth:`build`.
"""
return self.build(*args, **kwargs)
def build(self, inputs, **kwargs):
"""
Builds the LBN structure layer by layer within dedicated variable scopes. *input* must be a
tensor of the input four-vectors. All *kwargs* are forwarded to :py:meth:`build_features`.
"""
with tf.variable_scope(self.name):
with tf.variable_scope("placeholders"):
self.build_placeholders()
with tf.variable_scope("inputs"):
self.handle_input(inputs)
with tf.variable_scope("constants"):
self.build_constants()
with tf.variable_scope("particles"):
self.build_combinations("particle", self.n_particles)
# rest frames are not build for COMBINATIONS boost mode
if self.boost_mode != self.COMBINATIONS:
with tf.variable_scope("restframes"):
self.build_combinations("restframe", self.n_restframes)
with tf.variable_scope("boost"):
self.build_boost()
with tf.variable_scope("features"):
self.build_features(**kwargs)
self.features = self._raw_features
if self.batch_norm_center or self.batch_norm_scale:
with tf.variable_scope("norm"):
self.build_norm()
self.features = self._norm_features
return self.features
def build_placeholders(self):
"""
Builds the internal placeholders.
"""
# the train phase placeholder
if self.is_training is None:
self.is_training = tf.placeholder(tf.bool, name="is_training")
def handle_input(self, inputs):
"""
Takes the passed four-vector *inputs* and infers dimensions and some internal tensors.
"""
# store the input vectors
self.inputs = inputs
# infer sizes
self.n_in = self.inputs.shape[1].value
self.n_dim = self.inputs.shape[2].value
if self.n_dim != 4:
raise Exception("input dimension must be 4 to represent 4-vectors")
# split 4-vector components
names = ["E", "px", "py", "pz"]
split = [1, 1, 1, 1]
for t, name in zip(tf.split(self.inputs, split, axis=-1), names):
setattr(self, "inputs_" + name, tf.squeeze(t, -1))
def build_constants(self):
"""
Builds the internal constants for the boost matrix.
"""
# 4x4 identity
self.I = tf.constant(np.identity(4), tf.float32)
# U matrix
self.U = tf.constant([[-1, 0, 0, 0]] + 3 * [[0, -1, -1, -1]], tf.float32)
def build_combinations(self, name, m):
"""
Builds the combination layers which are quite similiar for particles and rest frames. Hence,
*name* must be either ``"particle"`` or ``"restframe"``, and *m* is the corresponding number
of combinations.
"""
if name not in ("particle", "restframe"):
raise ValueError("unknown combination name '{}'".format(name))
# determine the weight tensor shape
weight_shape = (self.n_in, m)
# name helper
name_ = lambda tmpl: tmpl.format(name)
# build the weight matrix, or check it if already set
W = getattr(self, name_("{}_weights"))
if W is None:
# build a new weight matrix
if isinstance(self.weight_init, tuple):
mean, stddev = self.weight_init
else:
mean, stddev = 0., 1. / m
W = tf.Variable(tf.random_normal(weight_shape, mean, stddev, dtype=tf.float32))
else:
# W is set externally, check the shape, consider batching
shape = tuple(W.shape.as_list())
if shape != weight_shape:
raise ValueError("external {}_weights shape {} does not match {}".format(
name, shape, weight_shape))
# store as raw weights before applying abs or clipping
W = tf.identity(W, name=name_("raw_{}_weights"))
# apply abs
if getattr(self, name_("abs_{}_weights")):
W = tf.abs(W, name=name_("abs_{}_weights"))
# apply clipping
clip = getattr(self, name_("clip_{}_weights"))
if clip is True:
clip = self.epsilon
if clip is not False:
W = tf.maximum(W, clip, name=name_("clipped_{}_weights"))
# assign a name to the final weights
W = tf.identity(W, name=name_("{}_weights"))
# create four-vectors of combinations
E = tf.matmul(self.inputs_E, W, name=name_("{}s_E"))
px = tf.matmul(self.inputs_px, W, name=name_("{}s_px"))
py = tf.matmul(self.inputs_py, W, name=name_("{}s_py"))
pz = tf.matmul(self.inputs_pz, W, name=name_("{}s_pz"))
# create the full 3- and 4-vector stacks again
p = tf.stack([px, py, pz], axis=-1, name=name_("{}s_pvec"))
q = tf.stack([E, px, py, pz], axis=-1, name=name_("{}s"))
# save all tensors for later inspection
setattr(self, name_("{}_weights"), W)
setattr(self, name_("{}s_E"), E)
setattr(self, name_("{}s_px"), px)
setattr(self, name_("{}s_py"), py)
setattr(self, name_("{}s_pz"), pz)
setattr(self, name_("{}s_pvec"), p)
setattr(self, name_("{}s"), q)
def build_boost(self):
"""
Builds the boosted particles depending on the requested boost mode. For infos on the boost
matrix, see `this link <https://en.wikipedia.org/wiki/Lorentz_transformation>`__. The
vectorized implementation is as follows:
I = identity(4x4)
U = -1(1x1) 0(1x3)
0(3x1) -1(3x3)
e = (1, -beta_vec/beta(1x3))^T
Lambda = I + (U + gamma) x ((U + 1) x beta - U) x e . e^T
"""
# n_particles and n_restframes must be identical for PAIRS and COMBINATIONS boosting
if self.boost_mode in (self.PAIRS, self.COMBINATIONS):
if self.n_restframes != self.n_particles:
raise ValueError("n_restframes ({}) must be identical to n_particles ({}) in boost"
" mode '{}'".format(self.n_restframes, self.n_particles, self.boost_mode))
# get the objects that are used to infer beta and gamma for the build the boost matrix,
if self.boost_mode == self.COMBINATIONS:
restframes_E = self.particles_E
restframes_pvec = self.particles_pvec
else:
restframes_E = self.restframes_E
restframes_pvec = self.restframes_pvec
# to build the boost parameters, reshape E and p tensors so that batch and particle axes
# are merged, and once the Lambda matrix is built, this reshape is reverted again
# note: there might be more performant operations in future TF releases
E = tf.reshape(restframes_E, [-1, 1])
pvec = tf.reshape(restframes_pvec, [-1, 3])
# determine the beta vectors
betavec = tf.div(pvec, E)
# determine the scalar beta and gamma values
beta = tf.div(tf.sqrt(tf.reduce_sum(tf.square(pvec), axis=1)), tf.squeeze(E, axis=-1))
gamma = tf.div(1., tf.sqrt(1. - tf.square(beta) + self.epsilon))
# the e vector, (1, -betavec / beta)^T
beta = tf.expand_dims(beta, axis=-1)
e = tf.expand_dims(tf.concat([tf.ones_like(E), -tf.div(betavec, beta)], axis=-1), axis=-1)
e_T = tf.transpose(e, perm=[0, 2, 1])
# finally, the boost matrix
beta = tf.expand_dims(beta, axis=-1)
gamma = tf.reshape(gamma, [-1, 1, 1])
Lambda = self.I + (self.U + gamma) * ((self.U + 1) * beta - self.U) * tf.matmul(e, e_T)
# revert the merging of batch and particle axes
Lambda = tf.reshape(Lambda, [-1, self.n_restframes, 4, 4])
# prepare particles for matmul
particles = tf.reshape(self.particles, [-1, self.n_particles, 4, 1])
# Lambda and particles need to be updated for PRODUCT and COMBINATIONS boosting
if self.boost_mode in (self.PRODUCT, self.COMBINATIONS):
# two approaches are possible
# a) tile Lambda while repeating particles
# b) batched gather using tiled and repeated indices
# go with b) for the moment since diagonal entries can be removed before the matmul
l_indices = np.tile(np.arange(self.n_restframes), self.n_particles)
p_indices = np.repeat(np.arange(self.n_particles), self.n_restframes)
# remove indices that would lead to diagonal entries for COMBINATIONS boosting
if self.boost_mode == self.COMBINATIONS:
no_diag = np.hstack((triu_range(self.n_particles), tril_range(self.n_particles)))
l_indices = l_indices[no_diag]
p_indices = p_indices[no_diag]
# update Lambda and particles
Lambda = tf.gather(Lambda, l_indices, axis=1)
particles = tf.gather(particles, p_indices, axis=1)
# store the final boost matrix
self.Lambda = Lambda
# actual boosting
boosted_particles = tf.matmul(self.Lambda, particles)
# remove the last dimension resulting from multiplication and save
self.boosted_particles = tf.squeeze(boosted_particles, axis=-1, name="boosted_particles")
def build_features(self, features=None, external_features=None):
"""
Builds the output features. *features* should be a list of feature names as registered to
the :py:attr:`feature_factory` instance. When *None*, the default features
``["E", "px", "py", "pz"]`` are build. *external_features* can be a list of tensors of
externally produced features, that are concatenated to the build features and are, e.g.,
subject to the internal batch normalization.
"""
# default to reshaped 4-vector elements
if features is None:
features = ["E", "px", "py", "pz"]
# create the list of feature ops to concat
concat = []
for name in features:
func = getattr(self.feature_factory, name)
if func is None:
raise ValueError("unknown feature '{}'".format(name))
concat.append(func())
# add external features
if external_features is not None:
if isinstance(external_features, (list, tuple)):
concat.extend(list(external_features))
else:
concat.append(external_features)
# save raw features
self._raw_features = tf.concat(concat, axis=-1)
def build_norm(self):
"""
Applies simple batch normalization with floating averages to the output features using
``tf.layers.batch_normalization``. Make sure to also run the operation returned by
``tf.get_collection(tf.GraphKeys.UPDATE_OPS)`` during each train step.
"""
self._norm_features = tf.layers.batch_normalization(
self.features,
axis=1,
training=self.is_training,
center=self.batch_norm_center,
scale=self.batch_norm_scale,
)
class FeatureFactoryBase(object):
"""
Base class of the feature factory. It does not implemented actual features but rather the
feature wrapping and tensor caching functionality. So-called hidden features are subject to
caching but are not supposed to be accessed by the LBN. They rather provide intermediate results
that are used in multiple places and retained for performance purposes.
"""
excluded_attributes = ["_wrap_feature", "_wrap_features", "lbn"]
def __init__(self, lbn):
super(FeatureFactoryBase, self).__init__()
# cached tensors stored by name
# contains also hidden features
self._tensor_cache = {}
# dict of registered, bound feature functions
# does not contain hidden features
self._feature_funcs = {}
# wrap all features defined in this class
self._wrap_features()
# reference to the lbn instance
self.lbn = lbn
# some shorthands
self.n = lbn.n_out
self.epsilon = lbn.epsilon
def _wrap_feature(self, func, name=None, hidden=None):
"""
Wraps and registers a feature function. It ensures that the stored function is bound to this
instance. *name* defaults to the actual function name. When *hidden* is *None*, the decision
is inferred from whether the feature name starts with an underscore.
"""
if not name:
name = func.__name__
if hidden is None:
hidden = name.startswith("_")