superpostion, labelmerger and datamerger docs

src/models/gan/discriminator.py

@@ -29,13 +29,15 @@ class BaseDiscriminator(tf.keras.Model):
             List of indices for particle distribution channels that should be
             generated. Defaults to None, which will use [0,1,2,3,4,5,6]
             (everything except nuclei).
+            See also src/models/gan/datamerger.py for defaults.
         ed_feature_list : list, optional
             List of indices for energy deposit channels that should be
             generated. Defaults to None, which will use [8] (nothing except
             sum/total energy deposit).
+            See also src/models/gan/datamerger.py for defaults.
         numparticle : int, optional
             Maximum number of particle types that should be generated (used
-            for onehot encoding of particle types).
+            for onehot encoding of particle types). Defaults to 6.
         **kwargs
             Additional keyword arguments that are passed down to the
             constructor of tf.keras.Model.
@@ -171,7 +173,7 @@ class BaseDiscriminator(tf.keras.Model):
             Index 2: energy deposit, layout: (batch, depth, channel)
             Index 3: particle distribution mask, layout: (batch, depth, channel)
             Index 4: energy deposit mask, layout: (batch, depth, channel)
-        training : bool
+        training : bool, optional
             Flag to indicate if the model is run in training mode or not.
             Defaults to False.
 

src/models/gan/generator.py

@@ -35,13 +35,15 @@ class BaseGenerator(tf.keras.Model):
             List of indices for particle distribution channels that should be
             generated. Defaults to None, which will use [0,1,2,3,4,5,6]
             (everything except nuclei).
+            See also src/models/gan/datamerger.py for defaults.
         ed_feature_list : list, optional
             List of indices for energy deposit channels that should be
             generated. Defaults to None, which will use [8] (nothing except
             sum/total energy deposit).
+            See also src/models/gan/datamerger.py for defaults.
         numparticle : int, optional
             Maximum number of particle types that should be generated (used
-            for onehot encoding of particle types).
+            for onehot encoding of particle types). Defaults to 6.
         **kwargs
             Additional keyword arguments that are passed down to the
             constructor of tf.keras.Model.
@@ -188,7 +190,7 @@ class BaseGenerator(tf.keras.Model):
         inputs : list
             Index 0: label, layout: (batch, features)
             Index 1-end: noise vectors, layout: (type, batch, dimensions)
-        training : bool
+        training : bool, optional
             Flag to indicate if the model is run in training mode or not.
             Defaults to False.
 

src/models/gan/loss/__init__.py

+"""Package that implements all loss functions."""
+
 from .wasserstein_distance import WassersteinDistance
 from .gradient_penalty import GradientPenalty
 

src/models/gan/loss/gradient_penalty.py

@@ -103,7 +103,7 @@ class GradientPenalty(tf.keras.Model):
             Index 2: real energy deposit, layout: (batch, depth, channel)
             Index 3: fake particle distribution, layout: (batch, depth, channel)
             Index 4: fake energy deposit, layout: (batch, depth, channel)
-        training : bool
+        training : bool, optional
             Flag to indicate if the model is run in training mode or not.
             Defaults to False.
 

src/models/gan/loss/wasserstein_distance.py

@@ -37,7 +37,7 @@ class WassersteinDistance(tf.keras.Model):
             Index 2: real energy deposit, layout: (batch, depth, channel)
             Index 3: fake particle distribution, layout: (batch, depth, channel)
             Index 4: fake energy deposit, layout: (batch, depth, channel)
-        training : bool
+        training : bool, optional
             Flag to indicate if the model is run in training mode or not.
             Defaults to False.
 

src/models/gan/utils/__init__.py

+"""Packages that implements all utility and multi purpose model functionality."""
+
 from .labelmerger import LabelMerger
 from .datanormalizer import DataNormalizer
 from .datanormalizer import DataDenormalizer

src/models/gan/utils/datamerger.py

@@ -18,8 +18,34 @@ def get_default_ed_feature_list():
 
 
 class DataMerger(tf.keras.layers.Layer):
+    """DataMerger handles joining particle distribution and energy deposit data.
+
+    Because internally, all profile types are treated similarly, this class
+    handles the merging of selected particle distribution and energy deposit
+    channels.
+
+    This class implements the tf.keras.layers.Layer interface.
+    """
 
     def __init__(self, pd_feature_list=None, ed_feature_list=None, **kwargs):
+        """Construct a DataMerger instance.
+
+        Parameters
+        ----------
+        pd_feature_list : list, optional
+            List of indices for particle distribution channels that should be
+            generated. Defaults to None, which will use [0,1,2,3,4,5,6]
+            (everything except nuclei).
+            See also src/models/gan/datamerger.py for defaults.
+        ed_feature_list : list, optional
+            List of indices for energy deposit channels that should be
+            generated. Defaults to None, which will use [8] (nothing except
+            sum/total energy deposit).
+            See also src/models/gan/datamerger.py for defaults.
+        **kwargs
+            Additional keyword arguments that are passed down to the
+            constructor of tf.keras.layers.Layer.
+        """
         super().__init__(**kwargs)
 
         if pd_feature_list is None:
@@ -37,6 +63,22 @@ class DataMerger(tf.keras.layers.Layer):
 
     @tf.function
     def call(self, inputs, training=False):
+        """TensorFlow model call function.
+
+        Parameters
+        ----------
+        inputs : list
+            Index 0: particle distribution, layout: (batch, depth, channel)
+            Index 0: energy deposit, layout: (batch, depth, channel)
+        training : bool, optional
+            Flag to indicate if the layer is run in training mode or not.
+            Defaults to False.
+
+        Returns
+        -------
+        tensor : tf.Tensor
+            Merged data tensor, layout (batch, depth, net-channels)
+        """
         particle_distribution = inputs[0]
         energy_deposit = inputs[1]
 
@@ -50,8 +92,36 @@ class DataMerger(tf.keras.layers.Layer):
 
 
 class DataSplitter(tf.keras.layers.Layer):
+    """DataSplitter handles splitting particle distribution and energy deposit data.
+
+    Because internally, all profile types are treated similarly, this class
+    handles the splitting of selected particle distribution and energy deposit
+    channels. This is the reverse operation performed by DataMerger.
+    Channels that are not present internally (because they are deselected)
+    will be filled with nans.
+
+    This class implements the tf.keras.layers.Layer interface.
+    """
 
     def __init__(self, pd_feature_list=None, ed_feature_list=None, **kwargs):
+        """Construct a DataSplitter instance.
+
+        Parameters
+        ----------
+        pd_feature_list : list, optional
+            List of indices for particle distribution channels that should be
+            generated. Defaults to None, which will use [0,1,2,3,4,5,6]
+            (everything except nuclei).
+            See also src/models/gan/datamerger.py for defaults.
+        ed_feature_list : list, optional
+            List of indices for energy deposit channels that should be
+            generated. Defaults to None, which will use [8] (nothing except
+            sum/total energy deposit).
+            See also src/models/gan/datamerger.py for defaults.
+        **kwargs
+            Additional keyword arguments that are passed down to the
+            constructor of tf.keras.layers.Layer.
+        """
         super().__init__(**kwargs)
 
         if pd_feature_list is None:
@@ -72,6 +142,22 @@ class DataSplitter(tf.keras.layers.Layer):
 
     @staticmethod
     def get_gatherindex(feature_list, numfeatures):
+        """Helper function that converts a feature list into tf.gather indices.
+        
+        Parameters
+        ----------
+        feature_list : list
+            Either a particle distribution or energy deposit list of channels.
+        numfeatures : int
+            The maximum number of features (channels) in the corresponding
+            particle distribution or energy deposit data.
+
+        Returns
+        -------
+        gatherindex : int
+            A list of indices used for tf.gather, which is used to select the
+            appropriate channels.
+        """
         gatherindex = np.arange(numfeatures) + len(feature_list)
         for ii,_ in enumerate(gatherindex):
             if ii in feature_list:
@@ -79,6 +165,16 @@ class DataSplitter(tf.keras.layers.Layer):
         return gatherindex
 
     def build(self, input_shape):
+        """TensorFlow model build function.
+
+        Performs length checks and constructs a dummy data instance, which
+        is filled by nans (this dummy will be masked with selected channels).
+
+        Parameters
+        ----------
+        input_shape
+            See tf.keras.layers.Layer.build.
+        """
         self._depthlen = input_shape[1]
 
         pd_numfeatures = len(self._pd_feature_list)
@@ -96,6 +192,25 @@ class DataSplitter(tf.keras.layers.Layer):
 
     @tf.function
     def call(self, inputs, training=False):
+        """TensorFlow model call function.
+
+        Parameters
+        ----------
+        inputs : tf.Tensor
+            A tensor that combines both particle distribution and
+            energy deposit data.
+        training : bool, optional
+            Flag to indicate if the layer is run in training mode or not.
+            Defaults to False.
+
+        Returns
+        -------
+        tensor : list
+            List of data tensors, layout: (batch, depth, channels)
+            Index 0: particle distribution
+            Index 1: energy deposit
+            Channels that are not present internally are filled with nans.
+        """
         splitindex = len(self._pd_feature_list)
         pd_part = inputs[:,:,:splitindex]
         ed_part = inputs[:,:,splitindex:]

src/models/gan/utils/labelmerger.py

@@ -2,8 +2,33 @@ import tensorflow as tf
 
 
 class LabelMerger(tf.keras.layers.Layer):
+    """This class is used to bring label data into the right format.
+
+    It performs onehot encoding of the particle type and logscaling for the
+    energy if exponential scaling is used for the model.
+    
+    This class implements the tf.keras.layers.Layer interface.
+    """
 
     def __init__(self, numparticle=6, expscale=False, epsilon=1e-25, **kwargs):
+        """Construct a LabelMerger instance.
+
+        Parameters
+        ----------
+        numparticle : int, optional
+            Maximum number of particle types that should be generated (used
+            for onehot encoding of particle types). Defaults to 6.
+        expscale : bool, optional
+            Flag to indicate if the model uses exponential scaling of outputs.
+            If so, the energy label should be logarithmically scaled.
+            Defaults to False.
+        epsilon : float, optional
+            Regularizing constant for the logscaling:
+                log(abs(x) + epsilon)
+        **kwargs
+            Additional keyword arguments that are passed down to the
+            constructor of tf.keras.layers.Layer.
+        """
         super().__init__(**kwargs)
         self._numparticle = numparticle
 
@@ -12,6 +37,21 @@ class LabelMerger(tf.keras.layers.Layer):
 
     @tf.function
     def call(self, inputs, training=False):
+        """TensorFlow model call function.
+
+        Parameters
+        ----------
+        inputs : tf.Tensor
+            Label input, layout: (batch, features)
+        training : bool, optional
+            Flag to indicate if the layer is run in training mode or not.
+            Defaults to False.
+
+        Returns
+        -------
+        tensor : tf.Tensor
+            Formated label output, layout: (batch, onehot + newfeatures)
+        """
         # normalize inputs
         particle = tf.cast(inputs[:,0], tf.int32)
         particle_oh = tf.one_hot(particle, self._numparticle, dtype=tf.float32)

src/models/gan/utils/superposition.py

@@ -3,11 +3,37 @@ import numpy as np
 
 
 class UniformSuperposition(tf.keras.layers.Layer):
+    """Build a superposition of real and fake data with random, uniform weights.
+
+    output = e * real + (1 - e) * fake
+    with e ~ U(0,1)
+
+    It is used for the gradient penalty calculation.
+    
+    This class implements the tf.keras.layers.Layer interface.
+    """
 
     def __init__(self, **kwargs):
+        """Construct a UniformSuperposition.
+
+        Parameters
+        ----------
+        **kwargs
+            Additional keyword arguments that are passed down to the
+            constructor of tf.keras.layers.Layer.
+        """
         super().__init__(**kwargs)
 
     def build(self, input_shape):
+        """TensorFlow model build function.
+
+        Performs shape checks and infers the shape of the random weights.
+
+        Parameters
+        ----------
+        input_shape
+            See tf.keras.layers.Layer.build.
+        """
         assert len(input_shape[0]) == len(input_shape[1])
         for ii in range(len(input_shape[0])):
             assert input_shape[0][ii] == input_shape[1][ii]
@@ -18,6 +44,22 @@ class UniformSuperposition(tf.keras.layers.Layer):
 
     @tf.function
     def call(self, inputs, training=False):
+        """TensorFlow model call function.
+
+        Parameters
+        ----------
+        inputs : list
+            Index 0: real data, layout: (batch, depth, channel)
+            Index 0: fake data, layout: (batch, depth, channel)
+        training : bool, optional
+            Flag to indicate if the layer is run in training mode or not.
+            Defaults to False.
+
+        Returns
+        -------
+        tensor : tf.Tensor
+            Superpostion tensor, layout: (batch, depth, channel)
+        """
         data0 = inputs[0]
         data1 = inputs[1]