From cc9fca00a946c8e64cbd3303b633cac6a5d50ca3 Mon Sep 17 00:00:00 2001
From: sascha-kirch <susch130993@googlemail.com>
Date: Fri, 20 Oct 2023 08:48:54 +0200
Subject: [PATCH 1/9] Add docstrings and refactor.

---
 DeepSaki/initializers/initializer_helper.py |   6 +-
 DeepSaki/layers/layer_composites.py         | 147 +++++++++++++-------
 DeepSaki/layers/layer_helper.py             |  23 +--
 DeepSaki/layers/padding.py                  |  40 +++++-
 DeepSaki/layers/pooling.py                  |   1 -
 DeepSaki/models/autoencoders.py             |   4 +-
 DeepSaki/models/discriminators.py           |  10 +-
 7 files changed, 159 insertions(+), 72 deletions(-)

diff --git a/DeepSaki/initializers/initializer_helper.py b/DeepSaki/initializers/initializer_helper.py
index df07045..b40e158 100644
--- a/DeepSaki/initializers/initializer_helper.py
+++ b/DeepSaki/initializers/initializer_helper.py
@@ -15,12 +15,14 @@ def make_initializer_complex(
     **Examples:**
     ```python
     # Standalone usage:
-    initializer = make_initializer_complex(tf.keras.initializers.GlorotUniform())
+    import DeepSaki as dsk
+    initializer = dsk.initializers.make_initializer_complex(tf.keras.initializers.GlorotUniform())
     values = initializer(shape=(2, 2))
     ```
     ```python
     # Usage in a Keras layer:
-    initializer = make_initializer_complex(tf.keras.initializers.GlorotUniform())
+    import DeepSaki as dsk
+    initializer = dsk.initializers.make_initializer_complex(tf.keras.initializers.GlorotUniform())
     layer = tf.keras.layers.Dense(3, kernel_initializer=initializer)
     ```
 
diff --git a/DeepSaki/layers/layer_composites.py b/DeepSaki/layers/layer_composites.py
index 265bbb4..c62d356 100644
--- a/DeepSaki/layers/layer_composites.py
+++ b/DeepSaki/layers/layer_composites.py
@@ -12,16 +12,12 @@
 from DeepSaki.layers.layer_helper import dropout_func
 from DeepSaki.layers.layer_helper import pad_func
 
+
 class Conv2DSplitted(tf.keras.layers.Layer):
-    """
-    To decrease the number of parameters, a convolution with the kernel_size (kernel,kernel) can be splitted into two consecutive convolutions with the kernel_size (kernel,1) and (1,kernel) respectivly
-    args:
-      - filters: number of filters in the output feature map
-      - kernels: size of the convolutions kernels, which will be translated to (kernels, 1) & (1,kernels) for the first and seccond convolution respectivly
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional  layers
-      - strides (optional, default: (1,1)): stride of the filter
-      - use_bias (optional, default: True): determines whether convolutions layers include a bias or not
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
+    """Convolution layer where a single convolution is splitted into two consecutive convolutions.
+
+    To decrease the number of parameters, a convolution with the kernel_size (kernel,kernel) can be splitted into two
+    consecutive convolutions with the kernel_size (kernel,1) and (1,kernel) respectivly.
     """
 
     def __init__(
@@ -33,6 +29,19 @@ def __init__(
         use_bias: bool = True,
         kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
     ) -> None:
+        """Initialize the `Conv2DSplitted` object.
+
+        Args:
+            filters (int): Number of filters in the output feature map
+            kernels (int): Size of the convolutions' kernels, which will be translated to (kernels, 1) and (1,kernels)
+                for the first and seccond convolution respectivly.
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers. Defaults
+                to False.
+            strides (Tuple[int, int], optional): Strides of the convolution layers. Defaults to (1, 1).
+            use_bias (bool, optional): Whether or not to use bias weights. Defaults to True.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions
+                kernels. Defaults to HeAlphaUniform().
+        """
         super(Conv2DSplitted, self).__init__()
         self.filters = filters
         self.kernels = kernels
@@ -61,6 +70,15 @@ def __init__(
             self.conv2 = tfa.layers.SpectralNormalization(self.conv2)
 
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
+        """Calls the `Conv2DSplitted` layer.
+
+        Args:
+            inputs (tf.Tensor): Tensor of shape `(batch, height, width, channel)`.
+
+        Returns:
+            Convoluted tensor of shape `(batch, H, W, filters)`, where `H` and `W` depend on the padding type used in
+                the convolution layers.
+        """
         x = self.conv1(inputs)
         return self.conv2(x)
 
@@ -149,22 +167,9 @@ def __init__(
                 )
                 for _ in range(number_of_convs)
             ]
-        else:
-            if use_spec_norm:
-                self.convs = [
-                    tfa.layers.SpectralNormalization(
-                        tf.keras.layers.Conv2D(
-                            filters=filters,
-                            kernel_size=(kernels, kernels),
-                            kernel_initializer=kernel_initializer,
-                            use_bias=use_bias,
-                            strides=strides,
-                        )
-                    )
-                    for _ in range(number_of_convs)
-                ]
-            else:
-                self.convs = [
+        elif use_spec_norm:
+            self.convs = [
+                tfa.layers.SpectralNormalization(
                     tf.keras.layers.Conv2D(
                         filters=filters,
                         kernel_size=(kernels, kernels),
@@ -172,13 +177,23 @@ def __init__(
                         use_bias=use_bias,
                         strides=strides,
                     )
-                    for _ in range(number_of_convs)
-                ]
-
-        if apply_final_normalization:
-            num_instancenorm_blocks = number_of_convs
+                )
+                for _ in range(number_of_convs)
+            ]
         else:
-            num_instancenorm_blocks = number_of_convs - 1
+            self.convs = [
+                tf.keras.layers.Conv2D(
+                    filters=filters,
+                    kernel_size=(kernels, kernels),
+                    kernel_initializer=kernel_initializer,
+                    use_bias=use_bias,
+                    strides=strides,
+                )
+                for _ in range(number_of_convs)
+            ]
+
+        num_instancenorm_blocks = number_of_convs if apply_final_normalization else number_of_convs - 1
+
         self.IN_blocks = [
             tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)
             for _ in range(num_instancenorm_blocks)
@@ -186,8 +201,12 @@ def __init__(
         self.dropout = dropout_func(filters, dropout_rate)
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(Conv2DBlock, self).build(input_shape)
-        # print("Model built with shape: {}".format(input_shape))
         self.residualConv = None
         if input_shape[-1] != self.filters and self.use_residual_Conv2DBlock:
             # split kernels for kernel_size = 1 not required
@@ -209,7 +228,7 @@ def call(self, inputs: tf.Tensor) -> tf.Tensor:
         for block in range(self.number_of_convs):
             residual = x
 
-            if self.pad != 0 and self.padding != "none":
+            if self.pad != 0 and self.padding != PaddingType.NONE:
                 x = pad_func(pad_values=(self.pad, self.pad), padding_type=self.padding)(x)
             x = self.convs[block](x)
 
@@ -272,7 +291,7 @@ class DenseBlock(tf.keras.layers.Layer):
     Wraps a dense layer into a more complex building block
     args:
       - units: number of units of each dense block
-      - numberOfLayers (optional, default: 1): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
+      - number_of_layers (optional, default: 1): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
       - activation (optional, default: "leaky_relu"): string literal to obtain activation function
       - dropout_rate (optional, default: 0): probability of the dropout layer. If the preceeding layer has more than one channel, spatial dropout is applied, otherwise standard dropout
       - final_activation (optional, default: True): whether or not to activate the output of this layer
@@ -286,7 +305,7 @@ class DenseBlock(tf.keras.layers.Layer):
     def __init__(
         self,
         units: int,
-        numberOfLayers: int = 1,
+        number_of_layers: int = 1,
         activation: str = "leaky_relu",
         dropout_rate: float = 0.0,
         final_activation: bool = True,
@@ -298,7 +317,7 @@ def __init__(
     ) -> None:
         super(DenseBlock, self).__init__()
         self.units = units
-        self.numberOfLayers = numberOfLayers
+        self.number_of_layers = number_of_layers
         self.dropout_rate = dropout_rate
         self.activation = activation
         self.final_activation = final_activation
@@ -313,15 +332,15 @@ def __init__(
                 tfa.layers.SpectralNormalization(
                     tf.keras.layers.Dense(units=units, use_bias=use_bias, kernel_initializer=kernel_initializer)
                 )
-                for _ in range(numberOfLayers)
+                for _ in range(number_of_layers)
             ]
         else:
             self.DenseBlocks = [
                 tf.keras.layers.Dense(units=units, use_bias=use_bias, kernel_initializer=kernel_initializer)
-                for _ in range(numberOfLayers)
+                for _ in range(number_of_layers)
             ]
 
-        num_instancenorm_blocks = numberOfLayers if apply_final_normalization else numberOfLayers - 1
+        num_instancenorm_blocks = number_of_layers if apply_final_normalization else number_of_layers - 1
         self.IN_blocks = [
             tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)
             for _ in range(num_instancenorm_blocks)
@@ -331,13 +350,13 @@ def __init__(
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
         x = inputs
 
-        for block in range(self.numberOfLayers):
+        for block in range(self.number_of_layers):
             x = self.DenseBlocks[block](x)
 
-            if block != (self.numberOfLayers - 1) or self.apply_final_normalization:
+            if block != (self.number_of_layers - 1) or self.apply_final_normalization:
                 x = self.IN_blocks[block](x)
 
-            if block != (self.numberOfLayers - 1) or self.final_activation:
+            if block != (self.number_of_layers - 1) or self.final_activation:
                 x = tf.keras.layers.Activation(self.activation)(x)
 
         if self.dropout_rate > 0:
@@ -354,7 +373,7 @@ def get_config(self) -> Dict[str, Any]:
         config.update(
             {
                 "units": self.units,
-                "numberOfLayers": self.numberOfLayers,
+                "number_of_layers": self.number_of_layers,
                 "dropout_rate": self.dropout_rate,
                 "activation": self.activation,
                 "final_activation": self.final_activation,
@@ -369,7 +388,7 @@ def get_config(self) -> Dict[str, Any]:
 
 
 # testcode
-# layer = DenseBlock(units = 512, numberOfLayers = 3, activation = "leaky_relu", apply_final_normalization=False)
+# layer = DenseBlock(units = 512, number_of_layers = 3, activation = "leaky_relu", apply_final_normalization=False)
 # print(layer.get_config())
 # DeepSaki.layers.plot_layer(layer,input_shape=(256,256,64))
 
@@ -410,6 +429,11 @@ def __init__(
         self.gamma_initializer = gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(DownSampleBlock, self).build(input_shape)
 
         self.layers = []
@@ -513,6 +537,11 @@ def __init__(
         self.padding = padding
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(UpSampleBlock, self).build(input_shape)
         self.layers = []
 
@@ -706,6 +735,11 @@ def __init__(
         self.dropout = dropout_func(filters, dropout_rate)
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(ResidualIdentityBlock, self).build(input_shape)
         self.conv0 = None
         if input_shape[-1] != self.filters:
@@ -734,7 +768,7 @@ def call(self, inputs: tf.Tensor) -> tf.Tensor:
         for block in range(self.number_of_blocks):
             residual = x
 
-            if self.pad != 0 and self.padding != "none":
+            if self.pad != 0 and self.padding != PaddingType.NONE:
                 x = pad_func(pad_values=(self.pad, self.pad), padding_type=self.padding)(x)
 
             # y is allways the footpoint of the cardinality blocks
@@ -818,6 +852,11 @@ def __init__(
         self.gamma_initializer = gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(ResBlockDown, self).build(input_shape)
 
         self.convRes = Conv2DBlock(
@@ -929,6 +968,11 @@ def __init__(
         self.gamma_initializer = gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(ResBlockUp, self).build(input_shape)
         self.convRes = Conv2DBlock(
             filters=input_shape[-1],
@@ -1060,6 +1104,11 @@ def __init__(
         self.gamma_initializer = gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(ScalarGatedSelfAttention, self).build(input_shape)
         batch_size, height, width, num_channel = input_shape
         if self.intermediateChannel is None:
@@ -1068,7 +1117,7 @@ def build(self, input_shape: tf.TensorShape) -> None:
         self.w_f = DenseBlock(
             units=self.intermediateChannel,
             use_spec_norm=self.use_spec_norm,
-            numberOfLayers=1,
+            number_of_layers=1,
             activation=None,
             apply_final_normalization=False,
             use_bias=False,
@@ -1078,7 +1127,7 @@ def build(self, input_shape: tf.TensorShape) -> None:
         self.w_g = DenseBlock(
             units=self.intermediateChannel,
             use_spec_norm=self.use_spec_norm,
-            numberOfLayers=1,
+            number_of_layers=1,
             activation=None,
             apply_final_normalization=False,
             use_bias=False,
@@ -1088,7 +1137,7 @@ def build(self, input_shape: tf.TensorShape) -> None:
         self.w_h = DenseBlock(
             units=self.intermediateChannel,
             use_spec_norm=self.use_spec_norm,
-            numberOfLayers=1,
+            number_of_layers=1,
             activation=None,
             apply_final_normalization=False,
             use_bias=False,
@@ -1098,7 +1147,7 @@ def build(self, input_shape: tf.TensorShape) -> None:
         self.w_fgh = DenseBlock(
             units=num_channel,
             use_spec_norm=self.use_spec_norm,
-            numberOfLayers=1,
+            number_of_layers=1,
             activation=None,
             apply_final_normalization=False,
             use_bias=False,
diff --git a/DeepSaki/layers/layer_helper.py b/DeepSaki/layers/layer_helper.py
index 2deb5de..26dcccc 100644
--- a/DeepSaki/layers/layer_helper.py
+++ b/DeepSaki/layers/layer_helper.py
@@ -97,11 +97,14 @@ def pad_func(
 
 
 def dropout_func(filters: int, dropout_rate: float) -> tf.keras.layers.Layer:
-    """
-    Wrapper to obtain a dropout layer depending on the size of the preceeding feature map
-    args:
-      - filters: number of filters of previous layer
-      - dropout_rate: probability with which dropout is performed
+    """Wrapper to obtain a dropout layer depending on the number of filters of the preceeding feature map.
+
+    Args:
+        filters (int): Number of filters of the preceeding layer.
+        dropout_rate (float): Probability of the dropout layer to drop weights.
+
+    Returns:
+        layer: Returns `tf.keras.layers.SpatialDropout2D` if number of filters > 1, otherwise 'tf.keras.layers.Dropout'.
     """
     if filters > 1:
         return tf.keras.layers.SpatialDropout2D(dropout_rate)
@@ -109,11 +112,11 @@ def dropout_func(filters: int, dropout_rate: float) -> tf.keras.layers.Layer:
 
 
 def plot_layer(layer: tf.keras.layers.Layer, input_shape: List[int]) -> None:
-    """
-    Creates an model from a given layer to be able to call model.summary() and to plot a graphic
-    args:
-      layer: tf.keras.layer object to be ploted
-      input_shape: shape of the input data without batchsize -> (height, width, channel)
+    """Creates a model from a given layer to be able to call model.summary() and to plot a graph image.
+
+    Args:
+        layer (tf.keras.layers.Layer): Layer to be plotted.
+        input_shape (List[int]): Shape of the desired input of the layer.
     """
     layer.build([None, *input_shape])
     inputs = tf.keras.layers.Input(shape=input_shape)
diff --git a/DeepSaki/layers/padding.py b/DeepSaki/layers/padding.py
index f1fee89..0f9e300 100644
--- a/DeepSaki/layers/padding.py
+++ b/DeepSaki/layers/padding.py
@@ -7,12 +7,29 @@
 import tensorflow as tf
 
 class ReflectionPadding2D(tf.keras.layers.Layer):
+    """Reflection Padding layer with support for TPU.
+    """
     def __init__(self, padding: Tuple[int, int] = (1, 1), **kwargs: Any) -> None:
+        """Initialize the `ReflectionPadding2D` layer.
+
+        Args:
+            padding (Tuple[int, int], optional): One-sided padding added to the `hight` and `width` to an input tensor
+                of shape (batch, height, width, channel)respectively. Defaults to (1, 1).
+        """
         super(ReflectionPadding2D, self).__init__(**kwargs)
         self.padding = tuple(padding)
 
     @tf.custom_gradient
     def _padding_func(self, input_tensor: tf.Tensor) -> Tuple[tf.Tensor, Callable[[tf.Tensor], tf.Tensor]]:
+        """Calculates results for forward prop and provides a function to calc. gradients for backward propagation.
+
+        Args:
+            input_tensor (tf.Tensor): Tensor of shape `(batch, height, width, channel)`.
+
+        Returns:
+            padded_tensor: Padded tensor of shape `(batch, height+2*padding, width+2*padding, channel)`.
+            custom_grad: Function returning the gradients during backprop.
+        """
         padding_width, padding_height = self.padding
         padding_tensor = [
             [0, 0],
@@ -20,7 +37,7 @@ def _padding_func(self, input_tensor: tf.Tensor) -> Tuple[tf.Tensor, Callable[[t
             [padding_width, padding_width],
             [0, 0],
         ]
-        result = tf.pad(input_tensor, padding_tensor, mode="REFLECT")
+        padded_tensor = tf.pad(input_tensor, padding_tensor, mode="REFLECT")
 
         # upstream gradient is the chainrule of all previous gradients!
         def custom_grad(upstream: tf.Tensor) -> tf.Tensor:
@@ -33,10 +50,19 @@ def custom_grad(upstream: tf.Tensor) -> tf.Tensor:
                 target_width=upstream.shape[2] - 2 * padding_width,
             )
 
-        return result, custom_grad
+        return padded_tensor, custom_grad
 
     def compute_output_shape(self, input_shape: tf.TensorShape) -> Tuple[int, int, int, int]:
-        """If you are using "channels_last" configuration"""
+        """Calculates the expected output shape after calling the layer.
+
+        Assumes "channels_last" configuration.
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input data to the layer.
+
+        Returns:
+            output_shape: expected output shape of the layer.
+        """
         return (
             input_shape[0],
             input_shape[1] + 2 * self.padding[0],
@@ -45,6 +71,14 @@ def compute_output_shape(self, input_shape: tf.TensorShape) -> Tuple[int, int, i
         )
 
     def call(self, input_tensor: tf.Tensor) -> tf.Tensor:
+        """Calls the `ReflectionPadding2D` layer.
+
+        Args:
+            input_tensor (tf.Tensor): Tensor of shape `(batch, height, width, channel)`.
+
+        Returns:
+            padded tensor of shape `(batch, height+2*padding, width+2*padding, channel)`.
+        """
         return self._padding_func(input_tensor)
 
     def get_config(self) -> Dict[str, Any]:
diff --git a/DeepSaki/layers/pooling.py b/DeepSaki/layers/pooling.py
index 1ed378b..9948b40 100644
--- a/DeepSaki/layers/pooling.py
+++ b/DeepSaki/layers/pooling.py
@@ -16,7 +16,6 @@ class FrequencyFilter(Enum):
         HIGH_PASS (int): Indicates that high frequency components shall be kept and low frequency components shall be
             filtered.
     """
-
     LOW_PASS = 1
     HIGH_PASS = 2
 
diff --git a/DeepSaki/models/autoencoders.py b/DeepSaki/models/autoencoders.py
index f7ad419..6b05eff 100644
--- a/DeepSaki/models/autoencoders.py
+++ b/DeepSaki/models/autoencoders.py
@@ -119,7 +119,7 @@ def __init__(self,
     self.bottleNeck = Bottleneck(use_ResidualIdentityBlock=use_ResidualIdentityBlock, n_bottleneck_blocks=n_bottleneck_blocks,use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation = activation, dropout_rate=dropout_rate, use_spec_norm=use_spec_norm, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     self.decoder = Decoder(number_of_levels=number_of_levels, upsampling=upsampling, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation=activation,dropout_rate=dropout_rate, use_ResidualIdentityBlock=use_ResidualIdentityBlock,use_spec_norm=use_spec_norm,use_self_attention=use_self_attention, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer,enable_skip_connections_input=True)
     if fully_connected == "MLP":
-      self.img_reconstruction = DenseBlock(units = input_shape[-1], use_spec_norm = use_spec_norm, numberOfLayers = 1, activation = final_activation, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+      self.img_reconstruction = DenseBlock(units = input_shape[-1], use_spec_norm = use_spec_norm, number_of_layers = 1, activation = final_activation, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     elif fully_connected == "1x1_conv":
       self.img_reconstruction = Conv2DBlock(filters = input_shape[-1],use_residual_Conv2DBlock = False, kernels = 1, split_kernels  = False, number_of_convs = 1, activation = final_activation, use_spec_norm=use_spec_norm, apply_final_normalization = False, use_bias = use_bias,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     #To enable mixed precission support for matplotlib and distributed training and to increase training stability
@@ -246,7 +246,7 @@ def __init__(self,
     self.bottleNeck = Bottleneck(use_ResidualIdentityBlock=use_ResidualIdentityBlock, n_bottleneck_blocks=n_bottleneck_blocks,use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs , activation = activation,dropout_rate=dropout_rate,use_spec_norm=use_spec_norm, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     self.decoder = Decoder(number_of_levels=number_of_levels, upsampling=upsampling, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation=activation,dropout_rate=dropout_rate, use_ResidualIdentityBlock=use_ResidualIdentityBlock,use_spec_norm=use_spec_norm,use_self_attention=use_self_attention, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     if fully_connected == "MLP":
-      self.img_reconstruction = DenseBlock(units = input_shape[-1], use_spec_norm = use_spec_norm, numberOfLayers = 1, activation = final_activation, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+      self.img_reconstruction = DenseBlock(units = input_shape[-1], use_spec_norm = use_spec_norm, number_of_layers = 1, activation = final_activation, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     elif fully_connected == "1x1_conv":
       self.img_reconstruction = Conv2DBlock(filters = input_shape[-1],use_residual_Conv2DBlock = False, kernels = 1, split_kernels  = False, number_of_convs = 1, activation = final_activation, use_spec_norm=use_spec_norm, apply_final_normalization = False,use_bias = use_bias,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     #To enable mixed precission support for matplotlib and distributed training and to increase training stability
diff --git a/DeepSaki/models/discriminators.py b/DeepSaki/models/discriminators.py
index 5d555f9..9414bba 100644
--- a/DeepSaki/models/discriminators.py
+++ b/DeepSaki/models/discriminators.py
@@ -101,9 +101,9 @@ def __init__(self,
       self.SA = None
 
     if fully_connected == "MLP":
-      self.cont1 = DenseBlock(units = filters * 8, use_spec_norm = use_spec_norm, numberOfLayers = number_of_convs, activation = activation, dropout_rate =dropout_rate, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
-      self.cont2 = DenseBlock(units = filters * 8, use_spec_norm = use_spec_norm, numberOfLayers = number_of_convs, activation = activation, dropout_rate =dropout_rate, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
-      self.cont3 = DenseBlock(units = filters * 8, use_spec_norm = use_spec_norm, numberOfLayers = number_of_convs, activation = activation, dropout_rate =0, final_activation=False, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+      self.cont1 = DenseBlock(units = filters * 8, use_spec_norm = use_spec_norm, number_of_layers = number_of_convs, activation = activation, dropout_rate =dropout_rate, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+      self.cont2 = DenseBlock(units = filters * 8, use_spec_norm = use_spec_norm, number_of_layers = number_of_convs, activation = activation, dropout_rate =dropout_rate, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+      self.cont3 = DenseBlock(units = filters * 8, use_spec_norm = use_spec_norm, number_of_layers = number_of_convs, activation = activation, dropout_rate =0, final_activation=False, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     elif fully_connected == "1x1_conv":
       self.cont1 = Conv2DBlock(filters=filters * 8, kernels = 1, activation = activation, split_kernels = split_kernels,number_of_convs=number_of_convs, use_residual_Conv2DBlock=False,dropout_rate=dropout_rate,use_spec_norm=use_spec_norm, padding=padding,use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
       self.cont2 = Conv2DBlock(filters=filters * 8, kernels = 1, activation = activation, split_kernels = split_kernels,number_of_convs=number_of_convs, use_residual_Conv2DBlock=False,dropout_rate=dropout_rate,use_spec_norm=use_spec_norm, padding=padding,use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
@@ -356,10 +356,10 @@ def __init__(self,
     self.bottleNeck = Bottleneck(use_ResidualIdentityBlock=use_ResidualIdentityBlock, n_bottleneck_blocks=n_bottleneck_blocks,use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs, activation = activation,dropout_rate=dropout_rate, channel_list=[16*ch], use_spec_norm=use_spec_norm, use_bias = use_bias,residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     self.decoder = Decoder(number_of_levels=number_of_levels, upsampling=upsampling, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation=activation,dropout_rate=dropout_rate, use_ResidualIdentityBlock=use_ResidualIdentityBlock, channel_list=[8*ch,8*ch,4*ch,2*ch,ch], use_spec_norm=use_spec_norm, use_bias = use_bias,residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer,enable_skip_connections_input=True)
     if fully_connected == "MLP":
-      self.img_reconstruction = DenseBlock(units = 1, use_spec_norm = use_spec_norm, numberOfLayers = 1, activation = None, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+      self.img_reconstruction = DenseBlock(units = 1, use_spec_norm = use_spec_norm, number_of_layers = 1, activation = None, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     elif fully_connected == "1x1_conv":
       self.img_reconstruction =Conv2DBlock(filters = 1, use_residual_Conv2DBlock = False, kernels = 1, split_kernels  = False, number_of_convs = 1, activation = None,use_spec_norm=use_spec_norm, apply_final_normalization = False, use_bias = use_bias,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
-    self.linear = DenseBlock(units = 1, use_spec_norm = use_spec_norm, numberOfLayers = 1, activation = None, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+    self.linear = DenseBlock(units = 1, use_spec_norm = use_spec_norm, number_of_layers = 1, activation = None, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     #To enable mixed precission support for matplotlib and distributed training and to increase training stability
     self.linear_dtype = tf.keras.layers.Activation("linear", dtype = tf.float32)
 

From 169f3b7b2fa0c98c61c28e603d01afeec3984358 Mon Sep 17 00:00:00 2001
From: sascha-kirch <susch130993@googlemail.com>
Date: Fri, 20 Oct 2023 10:21:13 +0200
Subject: [PATCH 2/9] add further doc strings

---
 DeepSaki/layers/layer_composites.py | 289 +++++++++++++++-------------
 1 file changed, 155 insertions(+), 134 deletions(-)

diff --git a/DeepSaki/layers/layer_composites.py b/DeepSaki/layers/layer_composites.py
index c62d356..976635f 100644
--- a/DeepSaki/layers/layer_composites.py
+++ b/DeepSaki/layers/layer_composites.py
@@ -103,25 +103,7 @@ def get_config(self) -> Dict[str, Any]:
 
 
 class Conv2DBlock(tf.keras.layers.Layer):
-    """
-    Wraps a two-dimensional convolution into a more complex building block
-    args:
-      - filters: number of filters in the output feature map
-      - kernels: size of the convolutions kernels
-      - use_residual_Conv2DBlock (optional, default: False):
-      - split_kernels (optional, default: False): to decrease the number of parameters, a convolution with the kernel_size (kernel,kernel) can be splitted into two consecutive convolutions with the kernel_size (kernel,1) and (1,kernel) respectivly
-      - number_of_convs (optional, default: 1): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
-      - activation (optional, default: "leaky_relu"): string literal to obtain activation function
-      - dropout_rate (optional, default: 0): probability of the dropout layer. If the preceeding layer has more than one channel, spatial dropout is applied, otherwise standard dropout
-      - final_activation (optional, default: True): whether or not to activate the output of this layer
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - strides (optional, default: (1,1)): stride of the filter
-      - padding (optional, default: "zero"): padding type. Options are "none", "zero" or "reflection"
-      - apply_final_normalization (optional, default: True): Whether or not to place a normalization on the layer's output
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
-    """
+    """Wraps a two-dimensional convolution into a more complex building block."""
 
     def __init__(
         self,
@@ -140,7 +122,36 @@ def __init__(
         use_bias: bool = True,
         kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
         gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-    ):
+    )->None:
+        """Initializes the `Conv2DBlock` layer.
+
+        Args:
+            filters (int): Number of individual filters.
+            kernels (int): Size of the convolutions kernels.
+            use_residual_Conv2DBlock (bool, optional): Adds a residual connection in parallel to the `Conv2DBlock`.
+                Defaults to False.
+            split_kernels (bool, optional): To decrease the number of parameters, a convolution with the kernel_size
+                `(kernel,kernel)` can be splitted into two consecutive convolutions with the kernel_size `(kernel,1)`
+                and `(1,kernel)` respectivly. Defaults to False.
+            number_of_convs (int, optional): Number of consecutive convolutional building blocks, i.e. `Conv2DBlock`.
+                Defaults to 1.
+            activation (str, optional): String literal or tensorflow activation function object to obtain activation
+                function. Defaults to "leaky_relu".
+            dropout_rate (float, optional): Probability of the dropout layer. If the preceeding layer has more than one
+                channel, spatial dropout is applied, otherwise standard dropout. Defaults to 0.0.
+            final_activation (bool, optional): Whether or not to activate the output of this layer. Defaults to True.
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers.
+                Defaults to False.
+            strides (Tuple[int, int], optional): Stride of the filter. Defaults to (1, 1).
+            padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
+            apply_final_normalization (bool, optional): Whether or not to place a normalization on the layer's output.
+                Defaults to True.
+            use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to HeAlphaUniform().
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to HeAlphaUniform().
+        """
         super(Conv2DBlock, self).__init__()
         self.filters = filters
         self.use_residual_Conv2DBlock = use_residual_Conv2DBlock
@@ -207,10 +218,10 @@ def build(self, input_shape: tf.TensorShape) -> None:
             input_shape (tf.TensorShape): Shape of the input tensor to this layer.
         """
         super(Conv2DBlock, self).build(input_shape)
-        self.residualConv = None
+        self.residual_conv = None
         if input_shape[-1] != self.filters and self.use_residual_Conv2DBlock:
             # split kernels for kernel_size = 1 not required
-            self.residualConv = tf.keras.layers.Conv2D(
+            self.residual_conv = tf.keras.layers.Conv2D(
                 filters=self.filters,
                 kernel_size=1,
                 kernel_initializer=self.kernel_initializer,
@@ -218,7 +229,7 @@ def build(self, input_shape: tf.TensorShape) -> None:
                 strides=self.strides,
             )
             if self.use_spec_norm:
-                self.residualConv = tfa.layers.SpectralNormalization(self.residualConv)
+                self.residual_conv = tfa.layers.SpectralNormalization(self.residual_conv)
 
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
         if not self.built:
@@ -234,9 +245,9 @@ def call(self, inputs: tf.Tensor) -> tf.Tensor:
 
             if self.use_residual_Conv2DBlock:
                 if (
-                    block == 0 and self.residualConv is not None
+                    block == 0 and self.residual_conv is not None
                 ):  # after the first conf, the channel depth matches between input and output
-                    residual = self.residualConv(residual)
+                    residual = self.residual_conv(residual)
                 x = tf.keras.layers.Add()([x, residual])
 
             if block != (self.number_of_convs - 1) or self.apply_final_normalization:
@@ -280,27 +291,8 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
-# testcode
-# layer = Conv2DBlock(filters = 128, use_residual_Conv2DBlock = True, kernels = 3, split_kernels  = True, use_spec_norm = True, number_of_convs = 3, activation = "relu", dropout_rate =0.1, final_activation = False, apply_final_normalization = True)
-# print(layer.get_config())
-# DeepSaki.layers.plot_layer(layer,input_shape=(256,256,64))
-
-
 class DenseBlock(tf.keras.layers.Layer):
-    """
-    Wraps a dense layer into a more complex building block
-    args:
-      - units: number of units of each dense block
-      - number_of_layers (optional, default: 1): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
-      - activation (optional, default: "leaky_relu"): string literal to obtain activation function
-      - dropout_rate (optional, default: 0): probability of the dropout layer. If the preceeding layer has more than one channel, spatial dropout is applied, otherwise standard dropout
-      - final_activation (optional, default: True): whether or not to activate the output of this layer
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - apply_final_normalization (optional, default: True): Whether or not to place a normalization on the layer's output
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
-    """
+    """Wraps a dense layer into a more complex building block."""
 
     def __init__(
         self,
@@ -315,6 +307,24 @@ def __init__(
         kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
         gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
     ) -> None:
+        """Initializes the `DenseBlock` layer.
+
+        Args:
+            units (int): Number of units of each dense block
+            number_of_layers (int, optional): Number of consecutive subblocks. Defaults to 1.
+            activation (str, optional): String literal to obtain activation function. Defaults to "leaky_relu".
+            dropout_rate (float, optional): Probability of the dropout layer dropping weights. If the preceeding layer
+                has more than one channel, spatial dropout is applied, otherwise standard dropout. Defaults to 0.0.
+            final_activation (bool, optional): Whether or not to activate the output of this layer. Defaults to True.
+            use_spec_norm (bool, optional): Applies spectral normalization to dense layers. Defaults to False.
+            apply_final_normalization (bool, optional): Whether or not to apply normalization at the layer's output.
+                Defaults to True.
+            use_bias (bool, optional): Whether dense layers include bias weights. Defaults to True.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to HeAlphaUniform().
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to HeAlphaUniform().
+        """
         super(DenseBlock, self).__init__()
         self.units = units
         self.number_of_layers = number_of_layers
@@ -386,26 +396,8 @@ def get_config(self) -> Dict[str, Any]:
         )
         return config
 
-
-# testcode
-# layer = DenseBlock(units = 512, number_of_layers = 3, activation = "leaky_relu", apply_final_normalization=False)
-# print(layer.get_config())
-# DeepSaki.layers.plot_layer(layer,input_shape=(256,256,64))
-
-
 class DownSampleBlock(tf.keras.layers.Layer):
-    """
-    Spatial down-sampling for grid-like data
-    args:
-      - downsampling (optional, default: "average_pooling"):
-      - activation (optional, default: "leaky_relu"): string literal to obtain activation function
-      - kernels (optional, default: 3): size of the convolution's kernels when using downsampling = "conv_stride_2"
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - padding (optional, default: "zero"): padding type. Options are "none", "zero" or "reflection"
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
-    """
+    """Spatial down-sampling for grid-like data using `Conv2DBlock`."""
 
     def __init__(
         self,
@@ -418,6 +410,22 @@ def __init__(
         kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
         gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
     ) -> None:
+        """Initializes the `DownSampleBlock` layer.
+
+        Args:
+            downsampling (str, optional): Downsampling method used. Defaults to "average_pooling".
+            activation (str, optional): String literal or tensorflow activation function object to obtain activation
+                function. Defaults to "leaky_relu".
+            kernels (int, optional): Size of the convolutions kernels. Defaults to 3.
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers. Defaults
+                to False.
+            padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
+            use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to HeAlphaUniform().
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to HeAlphaUniform().
+        """
         super(DownSampleBlock, self).__init__()
         self.kernels = kernels
         self.downsampling = downsampling
@@ -493,41 +501,38 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
-# layer = DownSampleBlock(numOfChannels = 3, kernels = 3, downsampling = "conv_stride_2", activation = "leaky_relu", use_spec_norm = True)
-# print(layer.get_config())
-# DeepSaki.layers.plot_layer(layer,input_shape=(256,256,64))
-
-
 class UpSampleBlock(tf.keras.layers.Layer):
-    """
-    Spatial up-sampling for grid-like data
-    args:
-      - upsampling (optional, default: "2D_upsample_and_conv"):
-      - activation (optional, default: "leaky_relu"): string literal to obtain activation function
-      - kernels (optional, default: 3): size of the convolution's kernels when using upsampling = "2D_upsample_and_conv" or "transpose_conv"
-      - split_kernels (optional, default: False): to decrease the number of parameters, a convolution with the kernel_size (kernel,kernel) can be splitted into two consecutive convolutions with the kernel_size (kernel,1) and (1,kernel) respectivly. applies to upsampling = "2D_upsample_and_conv"
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - padding (optional, default: "zero"): padding type. Options are "none", "zero" or "reflection"
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
-    """
+    """Spatial up-sampling for grid-like data using `Conv2DBlock`."""
 
     def __init__(
         self,
         upsampling: str = "2D_upsample_and_conv",
         activation: str = "leaky_relu",
         kernels: int = 3,
-        split_kernels: bool = False,
         use_spec_norm: bool = False,
         use_bias: bool = True,
         padding: PaddingType = PaddingType.ZERO,
         kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
         gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
     ) -> None:
+        """Initializes the `UpSampleBlock` layer.
+
+        Args:
+            upsampling (str, optional): _description_. Defaults to "2D_upsample_and_conv".
+            activation (str, optional): String literal or tensorflow activation function object to obtain activation
+                function. Defaults to "leaky_relu".
+            kernels (int, optional): Size of the convolutions kernels. Defaults to 3.
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers. Defaults
+                to False.
+            use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
+            padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to HeAlphaUniform().
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to HeAlphaUniform().
+        """
         super(UpSampleBlock, self).__init__()
         self.kernels = kernels
-        self.split_kernels = split_kernels
         self.activation = activation
         self.use_spec_norm = use_spec_norm
         self.upsampling = upsampling
@@ -552,7 +557,6 @@ def build(self, input_shape: tf.TensorShape) -> None:
                     filters=input_shape[-1],
                     use_residual_Conv2DBlock=False,
                     kernels=1,
-                    split_kernels=self.split_kernels,
                     number_of_convs=1,
                     activation=self.activation,
                     use_spec_norm=self.use_spec_norm,
@@ -625,11 +629,6 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
-# Testcode
-# layer = UpSampleBlock(kernels = 3, upsampling = "transpose_conv", activation = "leaky_relu", use_spec_norm = True, split_kernels = False)
-# print(layer.get_config())
-# DeepSaki.layers.plot_layer(layer,input_shape=(256,256,64))
-
 
 class ResidualIdentityBlock(tf.keras.layers.Layer):
     """
@@ -816,22 +815,27 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
-# Testcode
-# layer = ResidualIdentityBlock(filters =64, activation = "leaky_relu", kernels = 3, number_of_blocks=2,use_spec_norm = False, residual_cardinality =5, dropout_rate=0.2)
-# print(layer.get_config())
-# DeepSaki.layers.plot_layer(layer,input_shape=(256,256,64))
-
 
 class ResBlockDown(tf.keras.layers.Layer):
-    """
-    Spatial down-sampling with residual connection for grid-like data
-    args:
-      - activation (optional, default: "leaky_relu"): string literal to obtain activation function
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - padding (optional, default: "zero"): padding type. Options are "none", "zero" or "reflection"
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
+    """Spatial down-sampling with residual connection for grid-like data.
+
+    Architecture:
+        ```mermaid
+        flowchart LR
+            i([input_tensor])-->a1
+            i-->c2
+            subgraph ResBlockDown
+                subgraph Block[2x]
+                a1[activation]-->p1[padding]-->c1[Conv2DBlock]
+                end
+                c1-->ap1[AveragePooling2D]
+                c2[Conv2DBlock]-->ap2[AveragePooling2D]
+                ap1-->add((+))
+                ap2-->add((+))
+            end
+            add-->o([output_tensor])
+        ```
+
     """
 
     def __init__(
@@ -843,6 +847,20 @@ def __init__(
         kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
         gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
     ) -> None:
+        """Initializes the `ResBlockDown` layer.
+
+        Args:
+            activation (str, optional): String literal or tensorflow activation function object to obtain activation
+                function. Defaults to "leaky_relu".
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers. Defaults
+                to False.
+            use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
+            padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to HeAlphaUniform().
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to HeAlphaUniform().
+        """
         super(ResBlockDown, self).__init__()
         self.activation = activation
         self.use_spec_norm = use_spec_norm
@@ -932,12 +950,6 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
-# Testcode
-# layer = ResBlockDown( activation = "leaky_relu", use_spec_norm = True)
-# print(layer.get_config())
-# DeepSaki.layers.plot_layer(layer,input_shape=(256,256,64))
-
-
 class ResBlockUp(tf.keras.layers.Layer):
     """
     Spatial down-sampling with residual connection for grid-like data
@@ -1046,27 +1058,33 @@ def get_config(self) -> Dict[str, Any]:
         )
         return config
 
-
-# Testcode
-# layer = ResBlockUp(activation = "leaky_relu", use_spec_norm = True)
-# print(layer.get_config())
-# DeepSaki.layers.plot_layer(layer,input_shape=(256,256,64))
-
-
 @tf.keras.utils.register_keras_serializable(package="Custom", name="scale")
 class ScaleLayer(tf.keras.layers.Layer):
-    """
-    trainable scalar that can act as trainable gate
-    args:
-      - initializer (optional, default: tf.keras.initializers.Ones()): initializes the scalar weight
+    """Trainable non-negative scalar that might be used as a trainable gate.
+
+    It is a single learnable weight that is multiplied by all weights of the input tensor through broadcasting.
     """
 
     def __init__(self, initializer: tf.keras.initializers.Initializer = tf.keras.initializers.Ones()) -> None:
+        """Initializes the `ScaleLayer` layer.
+
+        Args:
+            initializer (tf.keras.initializers.Initializer, optional): Initializer used to initialize the scalar weight.
+                Defaults to tf.keras.initializers.Ones().
+        """
         super(ScaleLayer, self).__init__()
         self.initializer = initializer
         self.scale = self.add_weight(shape=[1], initializer=initializer, constraint=NonNegative(), trainable=True)
 
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
+        """Calls the `ScaleLayer` layer.
+
+        Args:
+            inputs (tf.Tensor): Tensor of shape (batch, ...).
+
+        Returns:
+            Non-negative scaled version of the input. Tensor of shape (batch, ...).
+        """
         return inputs * self.scale
 
     def get_config(self) -> Dict[str, Any]:
@@ -1081,13 +1099,10 @@ def get_config(self) -> Dict[str, Any]:
 
 
 class ScalarGatedSelfAttention(tf.keras.layers.Layer):
-    """
-    Scaled dot-product self attention that is gated by a learnable scalar.
-    args:
-    - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-    - intermediateChannel (optional, default: None): Integer that determines the intermediate channels within the self-attention model. If None, intermediate channels = inputChannels/8
-    - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-    - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
+    """Scaled dot-product self attention that is gated by a learnable scalar.
+
+    Info:
+        Implementation as used in the [VoloGAN paper](https://arxiv.org/abs/2207.09204).
     """
 
     def __init__(
@@ -1097,6 +1112,18 @@ def __init__(
         kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
         gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
     ) -> None:
+        """Initializes the `ScalarGatedSelfAttention` layer.
+
+        Args:
+            use_spec_norm (bool, optional): If true, applies spectral normalization to convolutional and dense layers.
+                Defaults to False.
+            intermediate_channel (Optional[bool], optional): Intermediate channels for the self attention mechanism.
+                Defaults to None.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the kernel weights.
+                Defaults to HeAlphaUniform().
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to HeAlphaUniform().
+        """
         super(ScalarGatedSelfAttention, self).__init__()
         self.use_spec_norm = use_spec_norm
         self.intermediateChannel = intermediate_channel
@@ -1201,9 +1228,3 @@ def get_config(self) -> Dict[str, Any]:
             }
         )
         return config
-
-
-# Testcode
-# layer =ScalarGatedSelfAttention(use_spec_norm = True,intermediateChannel = None)
-# print(layer.get_config())
-# DeepSaki.layers.plot_layer(layer,input_shape=(32,32,512))

From 350df1f87a3929aab099aea6be5c8e0fd99dff5b Mon Sep 17 00:00:00 2001
From: sascha-kirch <susch130993@googlemail.com>
Date: Fri, 20 Oct 2023 14:06:45 +0200
Subject: [PATCH 3/9] add docstrings. Rename ResnetIdentityBlock to ResnetBlock

---
 DeepSaki/layers/__init__.py             |   2 +-
 DeepSaki/layers/layer_composites.py     | 112 ++++++++++++++++--------
 DeepSaki/layers/sub_model_composites.py |  48 +++++-----
 DeepSaki/models/autoencoders.py         |  24 ++---
 DeepSaki/models/discriminators.py       |  12 +--
 5 files changed, 118 insertions(+), 80 deletions(-)

diff --git a/DeepSaki/layers/__init__.py b/DeepSaki/layers/__init__.py
index 89d708b..1c1cbbc 100644
--- a/DeepSaki/layers/__init__.py
+++ b/DeepSaki/layers/__init__.py
@@ -20,7 +20,7 @@
 from DeepSaki.layers.layer_composites import DenseBlock
 from DeepSaki.layers.layer_composites import DownSampleBlock
 from DeepSaki.layers.layer_composites import UpSampleBlock
-from DeepSaki.layers.layer_composites import ResidualIdentityBlock
+from DeepSaki.layers.layer_composites import ResidualBlock
 from DeepSaki.layers.layer_composites import ResBlockDown
 from DeepSaki.layers.layer_composites import ResBlockUp
 from DeepSaki.layers.layer_composites import ScaleLayer
diff --git a/DeepSaki/layers/layer_composites.py b/DeepSaki/layers/layer_composites.py
index 976635f..6bb9768 100644
--- a/DeepSaki/layers/layer_composites.py
+++ b/DeepSaki/layers/layer_composites.py
@@ -397,7 +397,7 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 class DownSampleBlock(tf.keras.layers.Layer):
-    """Spatial down-sampling for grid-like data using `Conv2DBlock`."""
+    """Spatial down-sampling for grid-like data using `DeepSaki.layers.ConvBlock2D()`."""
 
     def __init__(
         self,
@@ -502,7 +502,7 @@ def get_config(self) -> Dict[str, Any]:
 
 
 class UpSampleBlock(tf.keras.layers.Layer):
-    """Spatial up-sampling for grid-like data using `Conv2DBlock`."""
+    """Spatial up-sampling for grid-like data using `DeepSaki.layers.ConvBlock2D()`."""
 
     def __init__(
         self,
@@ -629,23 +629,9 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
+class ResidualBlock(tf.keras.layers.Layer):
+    """Residual block with configurable cardinality."""
 
-class ResidualIdentityBlock(tf.keras.layers.Layer):
-    """
-    Residual identity block with configurable cardinality
-    args:
-      - filters: number of filters in the output feature map
-      - kernels: size of the convolutions kernels
-      - number_of_blocks (optional, default: 1): number of consecutive convolutional building blocks.
-      - activation (optional, default: "leaky_relu"): string literal to obtain activation function
-      - dropout_rate (optional, default: 0): probability of the dropout layer. If the preceeding layer has more than one channel, spatial dropout is applied, otherwise standard dropout
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - residual_cardinality (optional, default: 1): number of parallel convolution blocks
-      - padding (optional, default: "zero"): padding type. Options are "none", "zero" or "reflection"
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
-    """
 
     def __init__(
         self,
@@ -661,7 +647,27 @@ def __init__(
         kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
         gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
     ) -> None:
-        super(ResidualIdentityBlock, self).__init__()
+        """Initializes the `ResidualBlock` layer.
+
+        Args:
+            filters (int): Number of individual filters.
+            kernels (int): Size of the convolutions kernels.
+            activation (str, optional): String literal or tensorflow activation function object to obtain activation
+                function. Defaults to "leaky_relu".
+            number_of_blocks (int, optional): Number of residual subblocks. Defaults to 1.
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers. Defaults
+                to False.
+            residual_cardinality (int, optional): Number of parallel paths. Defaults to 1.
+            dropout_rate (float, optional): Probability of the dropout layer. If the preceeding layer has more than one
+                channel, spatial dropout is applied, otherwise standard dropout. Defaults to 0.0.
+            use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
+            padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to HeAlphaUniform().
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to HeAlphaUniform().
+        """
+        super(ResidualBlock, self).__init__()
         self.activation = activation
         self.filters = filters
         self.kernels = kernels
@@ -739,7 +745,7 @@ def build(self, input_shape: tf.TensorShape) -> None:
         Args:
             input_shape (tf.TensorShape): Shape of the input tensor to this layer.
         """
-        super(ResidualIdentityBlock, self).build(input_shape)
+        super(ResidualBlock, self).build(input_shape)
         self.conv0 = None
         if input_shape[-1] != self.filters:
             self.conv0 = Conv2DBlock(
@@ -796,7 +802,7 @@ def get_config(self) -> Dict[str, Any]:
         Returns:
             Dictionary with the class' variable names as keys.
         """
-        config = super(ResidualIdentityBlock, self).get_config()
+        config = super(ResidualBlock, self).get_config()
         config.update(
             {
                 "activation": self.activation,
@@ -817,21 +823,19 @@ def get_config(self) -> Dict[str, Any]:
 
 
 class ResBlockDown(tf.keras.layers.Layer):
-    """Spatial down-sampling with residual connection for grid-like data.
+    """Spatial down-sampling with residual connection for grid-like data using `DeepSaki.layers.ConvBlock2D()`.
 
     Architecture:
         ```mermaid
         flowchart LR
-            i([input_tensor])-->a1
-            i-->c2
+            i([input_tensor])--> c1 & c2
             subgraph ResBlockDown
                 subgraph Block[2x]
-                a1[activation]-->p1[padding]-->c1[Conv2DBlock]
+                c1[dsk.layers.Conv2DBlock]
                 end
                 c1-->ap1[AveragePooling2D]
-                c2[Conv2DBlock]-->ap2[AveragePooling2D]
-                ap1-->add((+))
-                ap2-->add((+))
+                c2[dsk.layers.Conv2DBlock]-->ap2[AveragePooling2D]
+                ap1 & ap2-->add((+))
             end
             add-->o([output_tensor])
         ```
@@ -951,15 +955,23 @@ def get_config(self) -> Dict[str, Any]:
 
 
 class ResBlockUp(tf.keras.layers.Layer):
-    """
-    Spatial down-sampling with residual connection for grid-like data
-    args:
-      - activation (optional, default: "leaky_relu"): string literal to obtain activation function
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - padding (optional, default: "zero"): padding type. Options are "none", "zero" or "reflection"
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
+    """Spatial up-sampling with residual connection for grid-like data using `DeepSaki.layers.ConvBlock2D()`.
+
+    Architecture:
+        ```mermaid
+        flowchart LR
+            i([input_tensor])-->up1 & up2
+            subgraph ResBlockUp
+                up1[UpSample2D]-->c1
+                subgraph Block[2x]
+                c1[dsk.layers.Conv2DBlock]
+                end
+                up2[UpSample2D]-->c2[dsk.layers.Conv2DBlock]
+                c1 & c2-->add((+))
+            end
+            add-->o([output_tensor])
+        ```
+
     """
 
     def __init__(
@@ -971,6 +983,20 @@ def __init__(
         kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
         gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
     ) -> None:
+        """Initializes the `ResBlockUp` layer.
+
+        Args:
+            activation (str, optional): String literal or tensorflow activation function object to obtain activation
+                function. Defaults to "leaky_relu".
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers. Defaults
+                to False.
+            use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
+            padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to HeAlphaUniform().
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to HeAlphaUniform().
+        """
         super(ResBlockUp, self).__init__()
         self.activation = activation
         self.use_spec_norm = use_spec_norm
@@ -1103,6 +1129,18 @@ class ScalarGatedSelfAttention(tf.keras.layers.Layer):
 
     Info:
         Implementation as used in the [VoloGAN paper](https://arxiv.org/abs/2207.09204).
+
+    Architecture:
+        ```mermaid
+        flowchart LR
+            i([input_tensor])-->f & g & h
+            subgraph ScalarGatedSelfAttention
+                f[dsk.layers.DenseBlock] --> t[Transpose]
+                g[dsk.layers.DenseBlock] & t --> m1[Multiply] --> s[SoftMax]
+                h[dsk.layers.DenseBlock] & s --> m2[Multiply] --> v[DenseBlock] --> sc[dsk.layers.ScaleLayer]
+            end
+            sc -->o([output_tensor])
+        ```
     """
 
     def __init__(
diff --git a/DeepSaki/layers/sub_model_composites.py b/DeepSaki/layers/sub_model_composites.py
index 7dfa4ee..e23518e 100644
--- a/DeepSaki/layers/sub_model_composites.py
+++ b/DeepSaki/layers/sub_model_composites.py
@@ -12,7 +12,7 @@
 from DeepSaki.layers.layer_composites import DownSampleBlock
 from DeepSaki.layers.layer_composites import ResBlockDown
 from DeepSaki.layers.layer_composites import ResBlockUp
-from DeepSaki.layers.layer_composites import ResidualIdentityBlock
+from DeepSaki.layers.layer_composites import ResidualBlock
 from DeepSaki.layers.layer_composites import ScalarGatedSelfAttention
 from DeepSaki.layers.layer_composites import UpSampleBlock
 from DeepSaki.layers.layer_helper import PaddingType
@@ -32,8 +32,8 @@ class Encoder(tf.keras.layers.Layer):
       - number_of_convs (optional, default: 1): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
       - activation (optional, default: "leaky_relu"): string literal or tensorflow activation function object to obtain activation function
       - first_kernel (optional, default: 5): The first convolution can have a different kernel size, to e.g. increase the perceptive field, while the channel depth is still low.
-      - use_ResidualIdentityBlock (optional, default: False): Whether or not to use the ResidualIdentityBlock instead of the Conv2DBlock
-      - residual_cardinality (optional, default: 1): cardinality for the ResidualIdentityBlock
+      - use_ResidualBlock (optional, default: False): Whether or not to use the ResidualBlock instead of the Conv2DBlock
+      - residual_cardinality (optional, default: 1): cardinality for the ResidualBlock
       - channel_list (optional, default:None): alternativly to number_of_layers and filters, a list with the disired filters for each level can be provided. e.g. channel_list = [64, 128, 256] results in a 3-level Encoder with 64, 128, 256 filters for level 1, 2 and 3 respectivly.
       - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
       - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
@@ -58,7 +58,7 @@ def __init__(
         number_of_convs: int = 2,
         activation: str = "leaky_relu",
         first_kernel: Optional[int] = None,
-        use_ResidualIdentityBlock: bool = False,
+        use_ResidualBlock: bool = False,
         residual_cardinality: int = 1,
         channel_list: Optional[List[int]] = None,
         use_spec_norm: bool = False,
@@ -82,7 +82,7 @@ def __init__(
         self.number_of_convs = number_of_convs
         self.activation = activation
         self.first_kernel = first_kernel
-        self.use_ResidualIdentityBlock = use_ResidualIdentityBlock
+        self.use_ResidualBlock = use_ResidualBlock
         self.residual_cardinality = residual_cardinality
         self.channel_list = channel_list
         self.use_spec_norm = use_spec_norm
@@ -119,9 +119,9 @@ def build(self, input_shape: tf.TensorShape) -> None:
         for i, ch in enumerate(self.channel_list):
             encoder_kernels = self.first_kernel if i == 0 and self.first_kernel else self.kernels
 
-            if self.use_ResidualIdentityBlock:
+            if self.use_ResidualBlock:
                 self.encoderBlocks.append(
-                    ResidualIdentityBlock(
+                    ResidualBlock(
                         filters=ch,
                         activation=self.activation,
                         kernels=encoder_kernels,
@@ -216,7 +216,7 @@ def get_config(self) -> Dict[str, Any]:
                 "number_of_convs": self.number_of_convs,
                 "activation": self.activation,
                 "first_kernel": self.first_kernel,
-                "use_ResidualIdentityBlock": self.use_ResidualIdentityBlock,
+                "use_ResidualBlock": self.use_ResidualBlock,
                 "residual_cardinality": self.residual_cardinality,
                 "channel_list": self.channel_list,
                 "use_spec_norm": self.use_spec_norm,
@@ -234,7 +234,7 @@ def get_config(self) -> Dict[str, Any]:
 
 
 # Testcode
-# layer = Encoder( number_of_levels = 5, filters = 64, limit_filters = 512, use_self_attention = True,use_residual_Conv2DBlock = True, downsampling="max_pooling", kernels=3, split_kernels = True,  number_of_convs = 2,activation = "leaky_relu", first_kernel=3,use_ResidualIdentityBlock = True,use_spec_norm=True, omit_skips=2)
+# layer = Encoder( number_of_levels = 5, filters = 64, limit_filters = 512, use_self_attention = True,use_residual_Conv2DBlock = True, downsampling="max_pooling", kernels=3, split_kernels = True,  number_of_convs = 2,activation = "leaky_relu", first_kernel=3,use_ResidualBlock = True,use_spec_norm=True, omit_skips=2)
 # print(layer.get_config())
 # dsk.layers.plot_layer(layer,input_shape=(256,256,4))
 
@@ -248,13 +248,13 @@ class Bottleneck(tf.keras.layers.Layer):
       - split_kernels (optional, default: False): to decrease the number of parameters, a convolution with the kernel_size (kernel,kernel) can be splitted into two consecutive convolutions with the kernel_size (kernel,1) and (1,kernel) respectivly
       - number_of_convs (optional, default: 2): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
       - use_residual_Conv2DBlock (optional, default: True): ads a residual connection in parallel to the Conv2DBlock
-      - use_ResidualIdentityBlock (optional, default: False): Whether or not to use the ResidualIdentityBlock instead of the Conv2DBlock
+      - use_ResidualBlock (optional, default: False): Whether or not to use the ResidualBlock instead of the Conv2DBlock
       - activation (optional, default: "leaky_relu"): string literal or tensorflow activation function object to obtain activation function
       - dropout_rate (optional, default: 0): probability of the dropout layer. If the preceeding layer has more than one channel, spatial dropout is applied, otherwise standard dropout
       - channel_list (optional, default:None): alternativly to number_of_layers and filters, a list with the disired filters for each block can be provided. e.g. channel_list = [64, 128, 256] results in a 3-staged Bottleneck with 64, 128, 256 filters for stage 1, 2 and 3 respectivly.
       - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
       - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - residual_cardinality (optional, default: 1): cardinality for the ResidualIdentityBlock
+      - residual_cardinality (optional, default: 1): cardinality for the ResidualBlock
       - padding (optional, default: "none"): padding type. Options are "none", "zero" or "reflection"
       - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
       - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
@@ -267,7 +267,7 @@ def __init__(
         split_kernels: bool = False,
         number_of_convs: int = 2,
         use_residual_Conv2DBlock: bool = True,
-        use_ResidualIdentityBlock: bool = False,
+        use_ResidualBlock: bool = False,
         activation: str = "leaky_relu",
         dropout_rate: float = 0.2,
         channel_list: Optional[List[int]] = None,
@@ -279,7 +279,7 @@ def __init__(
         gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
     ) -> None:
         super(Bottleneck, self).__init__()
-        self.use_ResidualIdentityBlock = use_ResidualIdentityBlock
+        self.use_ResidualBlock = use_ResidualBlock
         self.n_bottleneck_blocks = n_bottleneck_blocks
         self.use_residual_Conv2DBlock = use_residual_Conv2DBlock
         self.kernels = kernels
@@ -304,9 +304,9 @@ def build(self, input_shape: tf.TensorShape) -> None:
 
         self.layers = []
         for ch in self.channel_list:
-            if self.use_ResidualIdentityBlock:
+            if self.use_ResidualBlock:
                 self.layers.append(
-                    ResidualIdentityBlock(
+                    ResidualBlock(
                         activation=self.activation,
                         filters=ch,
                         kernels=self.kernels,
@@ -360,7 +360,7 @@ def get_config(self) -> Dict[str, Any]:
         config = super(Bottleneck, self).get_config()
         config.update(
             {
-                "use_ResidualIdentityBlock": self.use_ResidualIdentityBlock,
+                "use_ResidualBlock": self.use_ResidualBlock,
                 "n_bottleneck_blocks": self.n_bottleneck_blocks,
                 "use_residual_Conv2DBlock": self.use_residual_Conv2DBlock,
                 "kernels": self.kernels,
@@ -400,8 +400,8 @@ class Decoder(tf.keras.layers.Layer):
       - number_of_convs (optional, default: 1): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
       - activation (optional, default: "leaky_relu"): string literal or tensorflow activation function object to obtain activation function
       - dropout_rate (optional, default: 0): probability of the dropout layer. If the preceeding layer has more than one channel, spatial dropout is applied, otherwise standard dropout. In the decoder only applied to the first half of levels.
-      - use_ResidualIdentityBlock (optional, default: False): Whether or not to use the ResidualIdentityBlock instead of the Conv2DBlock
-      - residual_cardinality (optional, default: 1): cardinality for the ResidualIdentityBlock
+      - use_ResidualBlock (optional, default: False): Whether or not to use the ResidualBlock instead of the Conv2DBlock
+      - residual_cardinality (optional, default: 1): cardinality for the ResidualBlock
       - channel_list (optional, default:None): alternativly to number_of_layers and filters, a list with the disired filters for each level can be provided. e.g. channel_list = [64, 128, 256] results in a 3-level Decoder with 64, 128, 256 filters for level 1, 2 and 3 respectivly.
       - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
       - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
@@ -424,7 +424,7 @@ def __init__(
         number_of_convs: int = 2,
         activation: str = "leaky_relu",
         dropout_rate: float = 0.2,
-        use_ResidualIdentityBlock: bool = False,
+        use_ResidualBlock: bool = False,
         residual_cardinality: int = 1,
         channel_list: Optional[List[int]] = None,
         use_spec_norm: bool = False,
@@ -445,7 +445,7 @@ def __init__(
         self.split_kernels = split_kernels
         self.number_of_convs = number_of_convs
         self.activation = activation
-        self.use_ResidualIdentityBlock = use_ResidualIdentityBlock
+        self.use_ResidualBlock = use_ResidualBlock
         self.channel_list = channel_list
         self.use_spec_norm = use_spec_norm
         self.use_bias = use_bias
@@ -484,9 +484,9 @@ def build(self, input_shape: tf.TensorShape) -> None:
         for i, ch in enumerate(self.channel_list):
             dropout_rate = self.dropout_rate if i < int(self.number_of_levels / 2) else 0
 
-            if self.use_ResidualIdentityBlock:
+            if self.use_ResidualBlock:
                 self.decoderBlocks.append(
-                    ResidualIdentityBlock(
+                    ResidualBlock(
                         filters=ch,
                         activation=self.activation,
                         kernels=self.kernels,
@@ -578,7 +578,7 @@ def get_config(self) -> Dict[str, Any]:
                 "split_kernels": self.split_kernels,
                 "number_of_convs": self.number_of_convs,
                 "activation": self.activation,
-                "use_ResidualIdentityBlock": self.use_ResidualIdentityBlock,
+                "use_ResidualBlock": self.use_ResidualBlock,
                 "residual_cardinality": self.residual_cardinality,
                 "channel_list": self.channel_list,
                 "use_spec_norm": self.use_spec_norm,
@@ -594,6 +594,6 @@ def get_config(self) -> Dict[str, Any]:
 
 
 # Testcode
-# layer = Decoder( number_of_levels = 5, filters = 64, limit_filters = 2048, use_self_attention = True,use_residual_Conv2DBlock = False, upsampling="depth_to_space", kernels=3, split_kernels = False,  number_of_convs = 2,activation = "leaky_relu",use_ResidualIdentityBlock = True,use_spec_norm=False, dropout_rate = 0.2)
+# layer = Decoder( number_of_levels = 5, filters = 64, limit_filters = 2048, use_self_attention = True,use_residual_Conv2DBlock = False, upsampling="depth_to_space", kernels=3, split_kernels = False,  number_of_convs = 2,activation = "leaky_relu",use_ResidualBlock = True,use_spec_norm=False, dropout_rate = 0.2)
 # print(layer.get_config())
 # dsk.layers.plot_layer(layer,input_shape=(256,256,4))
diff --git a/DeepSaki/models/autoencoders.py b/DeepSaki/models/autoencoders.py
index 6b05eff..56c42f5 100644
--- a/DeepSaki/models/autoencoders.py
+++ b/DeepSaki/models/autoencoders.py
@@ -54,7 +54,7 @@ def __init__(self,
             activation:str = "leaky_relu",
             limit_filters:int = 512,
             use_residual_Conv2DBlock:bool = False,
-            use_ResidualIdentityBlock:bool = False,
+            use_ResidualBlock:bool = False,
             residual_cardinality:int = 1,
             n_bottleneck_blocks:int = 1,
             dropout_rate:float = 0.2,
@@ -91,9 +91,9 @@ def __init__(self,
             Defaults to 512.
         use_residual_Conv2DBlock (bool, optional): Adds a residual connection in parallel to the `Conv2DBlock`. Defaults
             to False.
-        use_ResidualIdentityBlock (bool, optional): Whether or not to use the `ResidualIdentityBlock` instead of the
+        use_ResidualBlock (bool, optional): Whether or not to use the `ResidualBlock` instead of the
             `Conv2DBlock`. Defaults to False.
-        residual_cardinality (int, optional): Cardinality for the `ResidualIdentityBlock`. Defaults to 1.
+        residual_cardinality (int, optional): Cardinality for the `ResidualBlock`. Defaults to 1.
         n_bottleneck_blocks (int, optional): Number of consecutive convolution blocks in the bottleneck. Defaults to 1.
         dropout_rate (float, optional): Probability of the dropout layer. If the preceeding layer has more than one
             channel, spatial dropout is applied, otherwise standard dropout. Defaults to 0.2.
@@ -115,9 +115,9 @@ def __init__(self,
     """
     super(UNet, self).__init__()
 
-    self.encoder = Encoder(number_of_levels=number_of_levels, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, downsampling=downsampling, kernels=kernels, split_kernels=split_kernels, number_of_convs=number_of_convs,activation=activation, first_kernel=first_kernel,use_ResidualIdentityBlock=use_ResidualIdentityBlock,use_spec_norm=use_spec_norm, omit_skips=omit_skips, output_skips=True, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
-    self.bottleNeck = Bottleneck(use_ResidualIdentityBlock=use_ResidualIdentityBlock, n_bottleneck_blocks=n_bottleneck_blocks,use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation = activation, dropout_rate=dropout_rate, use_spec_norm=use_spec_norm, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
-    self.decoder = Decoder(number_of_levels=number_of_levels, upsampling=upsampling, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation=activation,dropout_rate=dropout_rate, use_ResidualIdentityBlock=use_ResidualIdentityBlock,use_spec_norm=use_spec_norm,use_self_attention=use_self_attention, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer,enable_skip_connections_input=True)
+    self.encoder = Encoder(number_of_levels=number_of_levels, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, downsampling=downsampling, kernels=kernels, split_kernels=split_kernels, number_of_convs=number_of_convs,activation=activation, first_kernel=first_kernel,use_ResidualBlock=use_ResidualBlock,use_spec_norm=use_spec_norm, omit_skips=omit_skips, output_skips=True, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+    self.bottleNeck = Bottleneck(use_ResidualBlock=use_ResidualBlock, n_bottleneck_blocks=n_bottleneck_blocks,use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation = activation, dropout_rate=dropout_rate, use_spec_norm=use_spec_norm, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+    self.decoder = Decoder(number_of_levels=number_of_levels, upsampling=upsampling, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation=activation,dropout_rate=dropout_rate, use_ResidualBlock=use_ResidualBlock,use_spec_norm=use_spec_norm,use_self_attention=use_self_attention, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer,enable_skip_connections_input=True)
     if fully_connected == "MLP":
       self.img_reconstruction = DenseBlock(units = input_shape[-1], use_spec_norm = use_spec_norm, number_of_layers = 1, activation = final_activation, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     elif fully_connected == "1x1_conv":
@@ -182,7 +182,7 @@ def __init__(self,
             activation:str = "leaky_relu",
             final_activation:str = "hard_sigmoid",
             use_residual_Conv2DBlock:bool = False,
-            use_ResidualIdentityBlock:bool = True,
+            use_ResidualBlock:bool = True,
             residual_cardinality:int = 32,
             limit_filters:int = 512,
             n_bottleneck_blocks:int = 5,
@@ -217,9 +217,9 @@ def __init__(self,
             function for the model's output activation. Defaults to "hard_sigmoid".
         use_residual_Conv2DBlock (bool, optional): Ads a residual connection in parallel to the `Conv2DBlock`. Defaults
             to False.
-        use_ResidualIdentityBlock (bool, optional): Whether or not to use the ResidualIdentityBlock instead of the
+        use_ResidualBlock (bool, optional): Whether or not to use the ResidualBlock instead of the
             `Conv2DBlock`. Defaults to False.
-        residual_cardinality (int, optional): Cardinality for the ResidualIdentityBlock. Defaults to 1.
+        residual_cardinality (int, optional): Cardinality for the ResidualBlock. Defaults to 1.
         limit_filters (int, optional): Limits the number of filters, which is doubled with every downsampling block.
             Defaults to 512.
         n_bottleneck_blocks (int, optional): Number of consecutive convolution blocks in the bottleneck. Defaults to 1.
@@ -242,9 +242,9 @@ def __init__(self,
     """
     super(ResNet, self).__init__()
 
-    self.encoder = Encoder(number_of_levels=number_of_levels, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, downsampling=downsampling, kernels=kernels, split_kernels=split_kernels, number_of_convs=number_of_convs,activation=activation, first_kernel=first_kernel,use_ResidualIdentityBlock=use_ResidualIdentityBlock,use_spec_norm=use_spec_norm, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
-    self.bottleNeck = Bottleneck(use_ResidualIdentityBlock=use_ResidualIdentityBlock, n_bottleneck_blocks=n_bottleneck_blocks,use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs , activation = activation,dropout_rate=dropout_rate,use_spec_norm=use_spec_norm, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
-    self.decoder = Decoder(number_of_levels=number_of_levels, upsampling=upsampling, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation=activation,dropout_rate=dropout_rate, use_ResidualIdentityBlock=use_ResidualIdentityBlock,use_spec_norm=use_spec_norm,use_self_attention=use_self_attention, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+    self.encoder = Encoder(number_of_levels=number_of_levels, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, downsampling=downsampling, kernels=kernels, split_kernels=split_kernels, number_of_convs=number_of_convs,activation=activation, first_kernel=first_kernel,use_ResidualBlock=use_ResidualBlock,use_spec_norm=use_spec_norm, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+    self.bottleNeck = Bottleneck(use_ResidualBlock=use_ResidualBlock, n_bottleneck_blocks=n_bottleneck_blocks,use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs , activation = activation,dropout_rate=dropout_rate,use_spec_norm=use_spec_norm, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+    self.decoder = Decoder(number_of_levels=number_of_levels, upsampling=upsampling, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation=activation,dropout_rate=dropout_rate, use_ResidualBlock=use_ResidualBlock,use_spec_norm=use_spec_norm,use_self_attention=use_self_attention, use_bias = use_bias, residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     if fully_connected == "MLP":
       self.img_reconstruction = DenseBlock(units = input_shape[-1], use_spec_norm = use_spec_norm, number_of_layers = 1, activation = final_activation, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     elif fully_connected == "1x1_conv":
diff --git a/DeepSaki/models/discriminators.py b/DeepSaki/models/discriminators.py
index 9414bba..e39ce9e 100644
--- a/DeepSaki/models/discriminators.py
+++ b/DeepSaki/models/discriminators.py
@@ -296,7 +296,7 @@ def __init__(self,
             number_of_convs:int = 2,
             activation:str = "leaky_relu",
             use_residual_Conv2DBlock:bool = False,
-            use_ResidualIdentityBlock:bool = False,
+            use_ResidualBlock:bool = False,
             residual_cardinality:int = 1,
             limit_filters:int = 512,
             n_bottleneck_blocks:int = 1,
@@ -328,9 +328,9 @@ def __init__(self,
             function. Defaults to "leaky_relu".
         use_residual_Conv2DBlock (bool, optional):Adds a residual connection in parallel to the `Conv2DBlock`. Defaults
             to False.
-        use_ResidualIdentityBlock (bool, optional): Whether or not to use the ResidualIdentityBlock instead of the
+        use_ResidualBlock (bool, optional): Whether or not to use the ResidualBlock instead of the
             `Conv2DBlock`. Defaults to False.
-        residual_cardinality (int, optional): Cardinality for the `ResidualIdentityBlock`. Defaults to 1.
+        residual_cardinality (int, optional): Cardinality for the `ResidualBlock`. Defaults to 1.
         limit_filters (int, optional): Limits the number of filters, which is doubled with every downsampling block.
             Defaults to 512.
         n_bottleneck_blocks (int, optional): Number of consecutive convolution blocks in the bottleneck. Defaults to 1.
@@ -352,9 +352,9 @@ def __init__(self,
     """
     super(UNetDiscriminator, self).__init__()
     ch = filters
-    self.encoder = Encoder(number_of_levels=number_of_levels, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, downsampling=downsampling, kernels=kernels, split_kernels=split_kernels, number_of_convs=number_of_convs,activation=activation, first_kernel=first_kernel, use_ResidualIdentityBlock=use_ResidualIdentityBlock, channel_list=[ch,2*ch,4*ch,8*ch,8*ch], use_spec_norm=use_spec_norm,use_self_attention=use_self_attention,output_skips=True, use_bias = use_bias,residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
-    self.bottleNeck = Bottleneck(use_ResidualIdentityBlock=use_ResidualIdentityBlock, n_bottleneck_blocks=n_bottleneck_blocks,use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs, activation = activation,dropout_rate=dropout_rate, channel_list=[16*ch], use_spec_norm=use_spec_norm, use_bias = use_bias,residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
-    self.decoder = Decoder(number_of_levels=number_of_levels, upsampling=upsampling, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation=activation,dropout_rate=dropout_rate, use_ResidualIdentityBlock=use_ResidualIdentityBlock, channel_list=[8*ch,8*ch,4*ch,2*ch,ch], use_spec_norm=use_spec_norm, use_bias = use_bias,residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer,enable_skip_connections_input=True)
+    self.encoder = Encoder(number_of_levels=number_of_levels, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, downsampling=downsampling, kernels=kernels, split_kernels=split_kernels, number_of_convs=number_of_convs,activation=activation, first_kernel=first_kernel, use_ResidualBlock=use_ResidualBlock, channel_list=[ch,2*ch,4*ch,8*ch,8*ch], use_spec_norm=use_spec_norm,use_self_attention=use_self_attention,output_skips=True, use_bias = use_bias,residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+    self.bottleNeck = Bottleneck(use_ResidualBlock=use_ResidualBlock, n_bottleneck_blocks=n_bottleneck_blocks,use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs, activation = activation,dropout_rate=dropout_rate, channel_list=[16*ch], use_spec_norm=use_spec_norm, use_bias = use_bias,residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
+    self.decoder = Decoder(number_of_levels=number_of_levels, upsampling=upsampling, filters=filters, limit_filters=limit_filters, use_residual_Conv2DBlock=use_residual_Conv2DBlock, kernels=kernels, split_kernels=split_kernels,number_of_convs=number_of_convs,activation=activation,dropout_rate=dropout_rate, use_ResidualBlock=use_ResidualBlock, channel_list=[8*ch,8*ch,4*ch,2*ch,ch], use_spec_norm=use_spec_norm, use_bias = use_bias,residual_cardinality=residual_cardinality,padding = padding, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer,enable_skip_connections_input=True)
     if fully_connected == "MLP":
       self.img_reconstruction = DenseBlock(units = 1, use_spec_norm = use_spec_norm, number_of_layers = 1, activation = None, apply_final_normalization = False, use_bias = use_bias, kernel_initializer = kernel_initializer, gamma_initializer = gamma_initializer)
     elif fully_connected == "1x1_conv":

From a47be4cd08747a0077fa6b3680e17e83f6ace867 Mon Sep 17 00:00:00 2001
From: sascha-kirch <susch130993@googlemail.com>
Date: Fri, 20 Oct 2023 14:55:41 +0200
Subject: [PATCH 4/9] added docstrings. Add iFFT3D layer

---
 DeepSaki/layers/fourier_layer.py | 160 +++++++++++++++++++++----------
 1 file changed, 111 insertions(+), 49 deletions(-)

diff --git a/DeepSaki/layers/fourier_layer.py b/DeepSaki/layers/fourier_layer.py
index 5392f7d..e09d8ff 100644
--- a/DeepSaki/layers/fourier_layer.py
+++ b/DeepSaki/layers/fourier_layer.py
@@ -372,28 +372,33 @@ def get_config(self) -> Dict[str, Any]:
         )
         return config
 
-
 class FFT2D(tf.keras.layers.Layer):
-    """
-    Calculates the 2D descrete fourier transform over the 2 innermost channels (height, width) over a 4 channel input (batch, height, width, channel)
-    args:
-    - is_channel_first: True or False. If True, input shape is assumed to be [batch,channel,height,width]. If False, input shape is assumed to be [batch,height,width,channel]
-    - applyRealFFT: True or False. If True, rfft2D is applied, which assumes real valued inputs and halves the width of the output. If False, fft2D is applied, which assumes complex input.
-    - shiftFFT: True or False. If true, low frequency componentes are centered.
-    - **kwargs: keyword arguments passed to the parent class tf.keras.layers.Layer.
+    """Calculates the 2D descrete fourier transform over the 2 innermost channels.
+
+    For a 4D input of shape (batch, height, width, channel), the 2DFFT would be calculated over (height, width)
     """
 
     def __init__(
         self,
         is_channel_first: bool = False,
-        applyRealFFT: bool = False,
-        shiftFFT: bool = True,
+        apply_real_fft: bool = False,
+        shift_fft: bool = True,
         **kwargs: Any,
     ) -> None:
+        """Initializes the `FFT2D` layer.
+
+        Args:
+            is_channel_first (bool, optional): Set true if input shape is (b,c,h,w) and false if input shape is
+                (b,h,w,c). Defaults to `False`.
+            apply_real_fft (bool, optional): If True, rfft2D is applied, which assumes real valued inputs and halves the
+                width of the output. If False, fft2D is applied, which assumes complex input. Defaults to False.
+            shift_fft (bool, optional): If true, low frequency componentes are centered. Defaults to True.
+            kwargs (Any): keyword arguments passed to the parent class tf.keras.layers.Layer.
+        """
         super(FFT2D, self).__init__(**kwargs)
         self.is_channel_first = is_channel_first
-        self.applyRealFFT = applyRealFFT
-        self.shiftFFT = shiftFFT
+        self.apply_real_fft = apply_real_fft
+        self.shift_fft = shift_fft
         self.policy_compute_dtype = tf.keras.mixed_precision.global_policy().compute_dtype
 
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
@@ -401,15 +406,15 @@ def call(self, inputs: tf.Tensor) -> tf.Tensor:
         if not self.is_channel_first:
             inputs = tf.einsum("bhwc->bchw", inputs)
 
-        if self.applyRealFFT:
+        if self.apply_real_fft:
             x = tf.signal.rfft2d(inputs)
-            if self.shiftFFT:
+            if self.shift_fft:
                 x = tf.signal.fftshift(x, axes=[-2])
         else:
             imag = tf.zeros_like(inputs)
             inputs = tf.complex(inputs, imag)  # fft2d requires complex inputs -> create complex with 0 imaginary
             x = tf.signal.fft2d(inputs)
-            if self.shiftFFT:
+            if self.shift_fft:
                 x = tf.signal.fftshift(x)
 
         if not self.is_channel_first:  # reverse the channel configuration to its initial config
@@ -426,8 +431,8 @@ def get_config(self) -> Dict[str, Any]:
         config.update(
             {
                 "is_channel_first": self.is_channel_first,
-                "applyRealFFT": self.applyRealFFT,
-                "shiftFFT": self.shiftFFT,
+                "apply_real_fft": self.apply_real_fft,
+                "shift_fft": self.shift_fft,
                 "policy_compute_dtype": self.policy_compute_dtype,
             }
         )
@@ -435,32 +440,41 @@ def get_config(self) -> Dict[str, Any]:
 
 
 class FFT3D(tf.keras.layers.Layer):
+    """Calculates the 3D descrete fourier transform over the 3 innermost channels.
+
+    For a 4D input like a batch of images of shape (batch, height, width, channel), the 3DFFT would be calculated over
+    (height, width).
+    For a 5D input, like a batch of videos of shape (batch, frame, height, width, channel), the 3DFFT would be
+    calculated over (frame, height, width).
+
     """
-    Calculates the 3D descrete fourier transform over the 3 innermost channels (height, width, channel) over a 4 channel input (batch, height, width, channel)
-    args:
-    - applyRealFFT: True or False. If True, rfft3D is applied, which assumes real valued inputs and halves the width of the output. If False, fft3D is applied, which assumes complex input.
-    - shiftFFT: True or False. If true, low frequency componentes are centered.
-    - **kwargs: keyword arguments passed to the parent class tf.keras.layers.Layer.
-    """
 
-    def __init__(self, applyRealFFT: bool = False, shiftFFT: bool = True, **kwargs: Any) -> None:
+    def __init__(self, apply_real_fft: bool = False, shift_fft: bool = True, **kwargs: Any) -> None:
+        """Initializes the `FFT3D` layer.
+
+        Args:
+            apply_real_fft (bool, optional): If True, rfft3D is applied, which assumes real valued inputs and halves the
+                width of the output. If False, fft3D is applied, which assumes complex input. Defaults to False.
+            shift_fft (bool, optional): If true, low frequency componentes are centered. Defaults to True.
+            kwargs (Any): keyword arguments passed to the parent class tf.keras.layers.Layer.
+        """
         super(FFT3D, self).__init__(**kwargs)
-        self.applyRealFFT = applyRealFFT
-        self.shiftFFT = shiftFFT
+        self.apply_real_fft = apply_real_fft
+        self.shift_fft = shift_fft
         self.policy_compute_dtype = tf.keras.mixed_precision.global_policy().compute_dtype
 
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
         inputs = tf.cast(inputs, tf.float32)  # mixed precision not supported with float16
 
-        if self.applyRealFFT:
+        if self.apply_real_fft:
             x = tf.signal.rfft3d(inputs)
-            if self.shiftFFT:
+            if self.shift_fft:
                 x = tf.signal.fftshift(x, axes=[-2])
         else:
             imag = tf.zeros_like(inputs)
             inputs = tf.complex(inputs, imag)  # fft3d requires complex inputs -> create complex with 0 imaginary
             x = tf.signal.fft3d(inputs)
-            if self.shiftFFT:
+            if self.shift_fft:
                 x = tf.signal.fftshift(x)
 
         return tf.cast(tf.math.real(x), self.policy_compute_dtype)
@@ -474,48 +488,49 @@ def get_config(self) -> Dict[str, Any]:
         config = super(FFT3D, self).get_config()
         config.update(
             {
-                "applyRealFFT": self.applyRealFFT,
-                "shiftFFT": self.shiftFFT,
+                "apply_real_fft": self.apply_real_fft,
+                "shift_fft": self.shift_fft,
                 "policy_compute_dtype": self.policy_compute_dtype,
             }
         )
         return config
 
-
 class iFFT2D(tf.keras.layers.Layer):
-    """
-    Calculates the 2D inverse FFT and reverses the center shift operation
-    args:
-    - is_channel_first: True or False. If True, input shape is assumed to be [batch,channel,height,width]. If False, input shape is assumed to be [batch,height,width,channel]
-    - applyRealFFT: True or False. If True, rfft2D is applied, which assumes real valued inputs and halves the width of the output. If False, fft2D is applied, which assumes complex input.
-    - shiftFFT: True or False. If True, shift operation of fourier transform is reversed before calculating the inverse fourier transformation
-    - **kwargs: keyword arguments passed to the parent class tf.keras.layers.Layer.
-
-    """
+    """Calculates the 2D inverse FFT."""
 
     def __init__(
         self,
         is_channel_first: bool = False,
-        applyRealFFT: bool = False,
-        shiftFFT: bool = True,
+        apply_real_fft: bool = False,
+        shift_fft: bool = True,
         **kwargs: Any,
     ) -> None:
+        """Initializes the `iFFT2D` layer.
+
+        Args:
+            is_channel_first (bool, optional): Set true if input shape is (b,c,h,w) and false if input shape is
+                (b,h,w,c). Defaults to `False`.
+            apply_real_fft (bool, optional): If True, rfft2D is applied, which assumes real valued inputs and halves the
+                width of the output. If False, fft2D is applied, which assumes complex input. Defaults to False.
+            shift_fft (bool, optional): If true, low frequency componentes are centered. Defaults to True.
+            kwargs (Any): keyword arguments passed to the parent class tf.keras.layers.Layer.
+        """
         super(iFFT2D, self).__init__(**kwargs)
         self.is_channel_first = is_channel_first
-        self.applyRealFFT = applyRealFFT
-        self.shiftFFT = shiftFFT
+        self.apply_real_fft = apply_real_fft
+        self.shift_fft = shift_fft
 
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
         if not self.is_channel_first:
             inputs = tf.einsum("bhwc->bchw", inputs)
         x = inputs
 
-        if self.applyRealFFT:
-            if self.shiftFFT:
+        if self.apply_real_fft:
+            if self.shift_fft:
                 x = tf.signal.ifftshift(x, axes=[-2])
             x = tf.signal.irfft2d(x)
         else:
-            if self.shiftFFT:
+            if self.shift_fft:
                 x = tf.signal.ifftshift(x)
             x = tf.signal.ifft2d(x)
 
@@ -533,6 +548,53 @@ def get_config(self) -> Dict[str, Any]:
         """
         config = super(iFFT2D, self).get_config()
         config.update(
-            {"is_channel_first": self.is_channel_first, "applyRealFFT": self.applyRealFFT, "shiftFFT": self.shiftFFT}
+            {"is_channel_first": self.is_channel_first, "apply_real_fft": self.apply_real_fft, "shift_fft": self.shift_fft}
+        )
+        return config
+
+class iFFT3D(tf.keras.layers.Layer):
+    """Calculates the 3D inverse FFT."""
+
+    def __init__(
+        self,
+        apply_real_fft: bool = False,
+        shift_fft: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes the `iFFT3D` layer.
+
+        Args:
+            apply_real_fft (bool, optional): If True, rfft3D is applied, which assumes real valued inputs and halves the
+                width of the output. If False, fft3D is applied, which assumes complex input. Defaults to False.
+            shift_fft (bool, optional): If true, low frequency componentes are centered. Defaults to True.
+            kwargs (Any): keyword arguments passed to the parent class tf.keras.layers.Layer.
+        """
+        super(iFFT3D, self).__init__(**kwargs)
+        self.apply_real_fft = apply_real_fft
+        self.shift_fft = shift_fft
+
+    def call(self, inputs: tf.Tensor) -> tf.Tensor:
+        x = inputs
+
+        if self.apply_real_fft:
+            if self.shift_fft:
+                x = tf.signal.ifftshift(x, axes=[-2])
+            x = tf.signal.irfft3d(x)
+        else:
+            if self.shift_fft:
+                x = tf.signal.ifftshift(x)
+            x = tf.signal.ifft3d(x)
+
+        return tf.math.real(x)
+
+    def get_config(self) -> Dict[str, Any]:
+        """Serialization of the object.
+
+        Returns:
+            Dictionary with the class' variable names as keys.
+        """
+        config = super(iFFT3D, self).get_config()
+        config.update(
+            {"apply_real_fft": self.apply_real_fft, "shift_fft": self.shift_fft}
         )
         return config

From 018611961385f332ff4fe27269925381b03e5468 Mon Sep 17 00:00:00 2001
From: sascha-kirch <susch130993@googlemail.com>
Date: Fri, 20 Oct 2023 16:49:35 +0200
Subject: [PATCH 5/9] add further doc strings

---
 .../activations/complex_valued_activations.py |   4 +-
 DeepSaki/initializers/he_alpha.py             |   2 +-
 DeepSaki/layers/fourier_layer.py              |  49 ++--
 DeepSaki/layers/layer_composites.py           | 137 +++++-----
 DeepSaki/layers/padding.py                    |   4 +-
 DeepSaki/layers/pooling.py                    |   1 +
 DeepSaki/layers/sub_model_composites.py       | 251 +++++++++++-------
 DeepSaki/models/autoencoders.py               |  16 +-
 DeepSaki/models/discriminators.py             |  18 +-
 9 files changed, 280 insertions(+), 202 deletions(-)

diff --git a/DeepSaki/activations/complex_valued_activations.py b/DeepSaki/activations/complex_valued_activations.py
index 338a827..c4b3298 100644
--- a/DeepSaki/activations/complex_valued_activations.py
+++ b/DeepSaki/activations/complex_valued_activations.py
@@ -13,11 +13,11 @@ class ComplexActivation(tf.keras.layers.Layer):
         tf.keras.layers.Layer
     """
 
-    def __init__(self, activation: tf.keras.layers.Layer = tf.keras.layers.ReLU(), **kwargs: Any) -> None:
+    def __init__(self, activation: tf.keras.layers.Layer, **kwargs: Any) -> None:
         """Initialize ComplexActivation.
 
         Args:
-            activation (tf.keras.layers.Layer, optional): Activation function to complexyfy. Defaults to tf.keras.layers.ReLU().
+            activation (tf.keras.layers.Layer): Activation function to complexyfy.
             kwargs: keyword arguments passed to the parent class tf.keras.layers.Layer.
         """
         super(ComplexActivation, self).__init__(**kwargs)
diff --git a/DeepSaki/initializers/he_alpha.py b/DeepSaki/initializers/he_alpha.py
index 0144fe0..3d758ec 100644
--- a/DeepSaki/initializers/he_alpha.py
+++ b/DeepSaki/initializers/he_alpha.py
@@ -99,7 +99,7 @@ def __init__(self, alpha: float = 0.3, seed: Optional[int] = None) -> None:
         """
         super(HeAlphaUniform, self).__init__(alpha, seed)
 
-    def __call__(self, shape: List[int], dtype: Optional[Union[tf.DType | np.dtype]] = None) -> tf.Tensor:
+    def __call__(self, shape: List[int], dtype: Optional[Union[tf.DType, np.dtype]] = None) -> tf.Tensor:
         """Dunder method to call the object instance.
 
         Args:
diff --git a/DeepSaki/layers/fourier_layer.py b/DeepSaki/layers/fourier_layer.py
index e09d8ff..82ec55a 100644
--- a/DeepSaki/layers/fourier_layer.py
+++ b/DeepSaki/layers/fourier_layer.py
@@ -29,8 +29,8 @@ def __init__(
         filters: int,
         kernels: Optional[Tuple[int, int]] = None,
         use_bias: bool = True,
-        kernel_initializer: tf.keras.initializers.Initializer = tf.keras.initializers.RandomUniform(-0.05, 0.05),
-        bias_initializer: tf.keras.initializers.Initializer = tf.keras.initializers.Zeros(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        bias_initializer: Optional[tf.keras.initializers.Initializer] = None,
         is_channel_first: bool = False,
         apply_conjugate: bool = False,
         pad_to_power_2: bool = True,
@@ -45,9 +45,9 @@ def __init__(
                 `[height,width]`. If `None`, kernel size is set to the input height and width. Defaults to `None`.
             use_bias (bool, optional): Whether or not to us bias weights. Defaults to `True`.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initializer to initialize the kernels of
-                the convolution layer. Defaults to `tf.keras.initializers.RandomUniform(-0.05, 0.05)`.
+                the convolution layer. Defaults to `None`.
             bias_initializer (tf.keras.initializers.Initializer, optional): Initializer to initialize the bias weights
-                of the convolution layer. Defaults to `tf.keras.initializers.Zeros()`.
+                of the convolution layer. Defaults to `None`.
             is_channel_first (bool, optional): Set true if input shape is (b,c,h,w) and false if input shape is
                 (b,h,w,c). Defaults to `False`.
             apply_conjugate (bool, optional): If true, the kernels are conjugated. If so, a multiplication in the
@@ -64,12 +64,15 @@ def __init__(
         self.filters = filters
         self.kernels = kernels
         self.use_bias = use_bias
-        self.kernel_initializer = kernel_initializer
-        self.bias_initializer = bias_initializer
+        self.kernel_initializer = (
+            tf.keras.initializers.RandomUniform(-0.05, 0.05) if kernel_initializer is None else kernel_initializer
+        )
+        self.bias_initializer = tf.keras.initializers.Zeros() if bias_initializer is None else bias_initializer
         self.is_channel_first = is_channel_first
         self.apply_conjugate = apply_conjugate
         self.pad_to_power_2 = pad_to_power_2
         self.method = method
+
         if method == MultiplicationType.MATRIX_PRODUCT:
             self.multiply = self._matrix_product
         elif method == MultiplicationType.ELEMENT_WISE:
@@ -257,8 +260,8 @@ def __init__(
         self,
         filters: int,
         use_bias: bool = True,
-        filter_initializer: tf.keras.initializers.Initializer = tf.keras.initializers.RandomUniform(-0.05, 0.05),
-        bias_initializer: tf.keras.initializers.Initializer = tf.keras.initializers.Zeros(),
+        filter_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        bias_initializer: Optional[tf.keras.initializers.Initializer] = None,
         is_channel_first: bool = False,
         **kwargs: Any,
     ) -> None:
@@ -268,9 +271,9 @@ def __init__(
             filters (int): Number of independent filters
             use_bias (bool, optional): Whether or not to us bias weights. Defaults to `True`.
             filter_initializer (tf.keras.initializers.Initializer, optional): Initializer to initialize the wheights of
-                the filter layer. Defaults to `tf.keras.initializers.RandomUniform(-0.05, 0.05)`.
+                the filter layer. Defaults to `None`.
             bias_initializer (tf.keras.initializers.Initializer, optional): Initializer to initialize the wheights of
-                the bias weights. Defaults to `tf.keras.initializers.Zeros()`.
+                the bias weights. Defaults to `None`.
             is_channel_first (bool, optional): Set true if input shape is (b,c,h,w) and false if input shape is
                 (b,h,w,c). Defaults to `False`.
             kwargs (Any): Additional key word arguments passed to the base class.
@@ -278,6 +281,11 @@ def __init__(
         super(FourierFilter2D, self).__init__(**kwargs)
         self.filters = filters
         self.use_bias = use_bias
+
+        filter_initializer = (
+            tf.keras.initializers.RandomUniform(-0.05, 0.05) if filter_initializer is None else filter_initializer
+        )
+        bias_initializer = tf.keras.initializers.Zeros() if bias_initializer is None else bias_initializer
         self.filter_initializer = make_initializer_complex(filter_initializer)
         self.bias_initializer = make_initializer_complex(bias_initializer)
         self.is_channel_first = is_channel_first
@@ -330,15 +338,13 @@ def call(self, inputs: tf.Tensor) -> tf.Tensor:
             inputs = tf.einsum("bhwc->bchw", inputs)
 
         output = np.ndarray(shape=self.out_shape)
-        for filter in range(self.filters):
+        for current_filter in range(self.filters):
+            # inputs:(batch, inp_filter, height, width ), fourier_filter:(...,inp_filter,height, width, out_filter)
             output = tf.concat(
                 [
                     output,
                     tf.reduce_sum(
-                        inputs
-                        * self.fourier_filter[
-                            :, :, :, filter
-                        ],  # inputs:(batch, inp_filter, height, width ), fourier_filter:(...,inp_filter,height, width, out_filter)
+                        inputs * self.fourier_filter[:, :, :, current_filter],
                         axis=-3,  # sum over all applied filters
                         keepdims=True,
                     ),
@@ -372,6 +378,7 @@ def get_config(self) -> Dict[str, Any]:
         )
         return config
 
+
 class FFT2D(tf.keras.layers.Layer):
     """Calculates the 2D descrete fourier transform over the 2 innermost channels.
 
@@ -495,6 +502,7 @@ def get_config(self) -> Dict[str, Any]:
         )
         return config
 
+
 class iFFT2D(tf.keras.layers.Layer):
     """Calculates the 2D inverse FFT."""
 
@@ -548,10 +556,15 @@ def get_config(self) -> Dict[str, Any]:
         """
         config = super(iFFT2D, self).get_config()
         config.update(
-            {"is_channel_first": self.is_channel_first, "apply_real_fft": self.apply_real_fft, "shift_fft": self.shift_fft}
+            {
+                "is_channel_first": self.is_channel_first,
+                "apply_real_fft": self.apply_real_fft,
+                "shift_fft": self.shift_fft,
+            }
         )
         return config
 
+
 class iFFT3D(tf.keras.layers.Layer):
     """Calculates the 3D inverse FFT."""
 
@@ -594,7 +607,5 @@ def get_config(self) -> Dict[str, Any]:
             Dictionary with the class' variable names as keys.
         """
         config = super(iFFT3D, self).get_config()
-        config.update(
-            {"apply_real_fft": self.apply_real_fft, "shift_fft": self.shift_fft}
-        )
+        config.update({"apply_real_fft": self.apply_real_fft, "shift_fft": self.shift_fft})
         return config
diff --git a/DeepSaki/layers/layer_composites.py b/DeepSaki/layers/layer_composites.py
index 6bb9768..85b13c5 100644
--- a/DeepSaki/layers/layer_composites.py
+++ b/DeepSaki/layers/layer_composites.py
@@ -12,7 +12,6 @@
 from DeepSaki.layers.layer_helper import dropout_func
 from DeepSaki.layers.layer_helper import pad_func
 
-
 class Conv2DSplitted(tf.keras.layers.Layer):
     """Convolution layer where a single convolution is splitted into two consecutive convolutions.
 
@@ -27,7 +26,7 @@ def __init__(
         use_spec_norm: bool = False,
         strides: Tuple[int, int] = (1, 1),
         use_bias: bool = True,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
         """Initialize the `Conv2DSplitted` object.
 
@@ -40,7 +39,7 @@ def __init__(
             strides (Tuple[int, int], optional): Strides of the convolution layers. Defaults to (1, 1).
             use_bias (bool, optional): Whether or not to use bias weights. Defaults to True.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions
-                kernels. Defaults to HeAlphaUniform().
+                kernels. Defaults to None.
         """
         super(Conv2DSplitted, self).__init__()
         self.filters = filters
@@ -48,7 +47,7 @@ def __init__(
         self.use_spec_norm = use_spec_norm
         self.strides = strides
         self.use_bias = use_bias
-        self.kernel_initializer = kernel_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
 
         self.conv1 = tf.keras.layers.Conv2D(
             filters=filters,
@@ -120,9 +119,9 @@ def __init__(
         padding: PaddingType = PaddingType.ZERO,
         apply_final_normalization: bool = True,
         use_bias: bool = True,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-    )->None:
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
+    ) -> None:
         """Initializes the `Conv2DBlock` layer.
 
         Args:
@@ -148,9 +147,9 @@ def __init__(
                 Defaults to True.
             use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
             gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
         """
         super(Conv2DBlock, self).__init__()
         self.filters = filters
@@ -166,8 +165,8 @@ def __init__(
         self.padding = padding
         self.apply_final_normalization = apply_final_normalization
         self.use_bias = use_bias
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
 
         self.pad = int((kernels - 1) / 2)  # assumes odd kernel size, which is typical!
 
@@ -184,7 +183,7 @@ def __init__(
                     tf.keras.layers.Conv2D(
                         filters=filters,
                         kernel_size=(kernels, kernels),
-                        kernel_initializer=kernel_initializer,
+                        kernel_initializer=self.kernel_initializer,
                         use_bias=use_bias,
                         strides=strides,
                     )
@@ -196,7 +195,7 @@ def __init__(
                 tf.keras.layers.Conv2D(
                     filters=filters,
                     kernel_size=(kernels, kernels),
-                    kernel_initializer=kernel_initializer,
+                    kernel_initializer=self.kernel_initializer,
                     use_bias=use_bias,
                     strides=strides,
                 )
@@ -304,8 +303,8 @@ def __init__(
         use_spec_norm: bool = False,
         apply_final_normalization: bool = True,
         use_bias: bool = True,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
         """Initializes the `DenseBlock` layer.
 
@@ -321,9 +320,9 @@ def __init__(
                 Defaults to True.
             use_bias (bool, optional): Whether dense layers include bias weights. Defaults to True.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
             gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
         """
         super(DenseBlock, self).__init__()
         self.units = units
@@ -334,25 +333,25 @@ def __init__(
         self.use_spec_norm = use_spec_norm
         self.apply_final_normalization = apply_final_normalization
         self.use_bias = use_bias
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
 
         if use_spec_norm:
             self.DenseBlocks = [
                 tfa.layers.SpectralNormalization(
-                    tf.keras.layers.Dense(units=units, use_bias=use_bias, kernel_initializer=kernel_initializer)
+                    tf.keras.layers.Dense(units=units, use_bias=use_bias, kernel_initializer=self.kernel_initializer)
                 )
                 for _ in range(number_of_layers)
             ]
         else:
             self.DenseBlocks = [
-                tf.keras.layers.Dense(units=units, use_bias=use_bias, kernel_initializer=kernel_initializer)
+                tf.keras.layers.Dense(units=units, use_bias=use_bias, kernel_initializer=self.kernel_initializer)
                 for _ in range(number_of_layers)
             ]
 
         num_instancenorm_blocks = number_of_layers if apply_final_normalization else number_of_layers - 1
         self.IN_blocks = [
-            tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)
+            tfa.layers.InstanceNormalization(gamma_initializer=self.gamma_initializer)
             for _ in range(num_instancenorm_blocks)
         ]
         self.dropout = tf.keras.layers.Dropout(dropout_rate)
@@ -396,6 +395,7 @@ def get_config(self) -> Dict[str, Any]:
         )
         return config
 
+
 class DownSampleBlock(tf.keras.layers.Layer):
     """Spatial down-sampling for grid-like data using `DeepSaki.layers.ConvBlock2D()`."""
 
@@ -407,8 +407,8 @@ def __init__(
         use_spec_norm: bool = False,
         padding: PaddingType = PaddingType.ZERO,
         use_bias: bool = True,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
         """Initializes the `DownSampleBlock` layer.
 
@@ -422,9 +422,9 @@ def __init__(
             padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
             use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
             gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
         """
         super(DownSampleBlock, self).__init__()
         self.kernels = kernels
@@ -433,8 +433,8 @@ def __init__(
         self.use_spec_norm = use_spec_norm
         self.padding = padding
         self.use_bias = use_bias
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
         """Build layer depending on the `input_shape` (output shape of the previous layer).
@@ -512,8 +512,8 @@ def __init__(
         use_spec_norm: bool = False,
         use_bias: bool = True,
         padding: PaddingType = PaddingType.ZERO,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
         """Initializes the `UpSampleBlock` layer.
 
@@ -527,9 +527,9 @@ def __init__(
             use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
             padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
             gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
         """
         super(UpSampleBlock, self).__init__()
         self.kernels = kernels
@@ -537,8 +537,8 @@ def __init__(
         self.use_spec_norm = use_spec_norm
         self.upsampling = upsampling
         self.use_bias = use_bias
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
         self.padding = padding
 
     def build(self, input_shape: tf.TensorShape) -> None:
@@ -632,7 +632,6 @@ def get_config(self) -> Dict[str, Any]:
 class ResidualBlock(tf.keras.layers.Layer):
     """Residual block with configurable cardinality."""
 
-
     def __init__(
         self,
         filters: int,
@@ -644,8 +643,8 @@ def __init__(
         dropout_rate: float = 0.0,
         use_bias: bool = True,
         padding: PaddingType = PaddingType.ZERO,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
         """Initializes the `ResidualBlock` layer.
 
@@ -663,9 +662,9 @@ def __init__(
             use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
             padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
             gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
         """
         super(ResidualBlock, self).__init__()
         self.activation = activation
@@ -677,8 +676,8 @@ def __init__(
         self.dropout_rate = dropout_rate
         self.use_bias = use_bias
         self.padding = padding
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
 
         self.pad = int((kernels - 1) / 2)  # assumes odd kernel size, which is typical!
 
@@ -704,8 +703,8 @@ def __init__(
                             use_spec_norm=use_spec_norm,
                             use_bias=use_bias,
                             padding=padding,
-                            kernel_initializer=kernel_initializer,
-                            gamma_initializer=gamma_initializer,
+                            kernel_initializer=self.kernel_initializer,
+                            gamma_initializer=self.gamma_initializer,
                         ),
                         Conv2DBlock(
                             filters=self.intermediateFilters,
@@ -717,8 +716,8 @@ def __init__(
                             padding=PaddingType.NONE,
                             use_spec_norm=use_spec_norm,
                             use_bias=use_bias,
-                            kernel_initializer=kernel_initializer,
-                            gamma_initializer=gamma_initializer,
+                            kernel_initializer=self.kernel_initializer,
+                            gamma_initializer=self.gamma_initializer,
                         ),
                         Conv2DBlock(
                             filters=filters,
@@ -730,8 +729,8 @@ def __init__(
                             use_spec_norm=use_spec_norm,
                             use_bias=use_bias,
                             padding=padding,
-                            kernel_initializer=kernel_initializer,
-                            gamma_initializer=gamma_initializer,
+                            kernel_initializer=self.kernel_initializer,
+                            gamma_initializer=self.gamma_initializer,
                         ),
                     ]
                 )
@@ -821,7 +820,6 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
-
 class ResBlockDown(tf.keras.layers.Layer):
     """Spatial down-sampling with residual connection for grid-like data using `DeepSaki.layers.ConvBlock2D()`.
 
@@ -848,8 +846,8 @@ def __init__(
         use_spec_norm: bool = False,
         use_bias: bool = True,
         padding: PaddingType = PaddingType.ZERO,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
         """Initializes the `ResBlockDown` layer.
 
@@ -861,17 +859,17 @@ def __init__(
             use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
             padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
             gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
         """
         super(ResBlockDown, self).__init__()
         self.activation = activation
         self.use_spec_norm = use_spec_norm
         self.use_bias = use_bias
         self.padding = padding
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
         """Build layer depending on the `input_shape` (output shape of the previous layer).
@@ -980,8 +978,8 @@ def __init__(
         use_spec_norm: bool = False,
         use_bias: bool = True,
         padding: PaddingType = PaddingType.ZERO,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
         """Initializes the `ResBlockUp` layer.
 
@@ -993,17 +991,17 @@ def __init__(
             use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
             padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
             gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
         """
         super(ResBlockUp, self).__init__()
         self.activation = activation
         self.use_spec_norm = use_spec_norm
         self.use_bias = use_bias
         self.padding = padding
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
         """Build layer depending on the `input_shape` (output shape of the previous layer).
@@ -1084,6 +1082,7 @@ def get_config(self) -> Dict[str, Any]:
         )
         return config
 
+
 @tf.keras.utils.register_keras_serializable(package="Custom", name="scale")
 class ScaleLayer(tf.keras.layers.Layer):
     """Trainable non-negative scalar that might be used as a trainable gate.
@@ -1091,15 +1090,15 @@ class ScaleLayer(tf.keras.layers.Layer):
     It is a single learnable weight that is multiplied by all weights of the input tensor through broadcasting.
     """
 
-    def __init__(self, initializer: tf.keras.initializers.Initializer = tf.keras.initializers.Ones()) -> None:
+    def __init__(self, initializer: Optional[tf.keras.initializers.Initializer] = None) -> None:
         """Initializes the `ScaleLayer` layer.
 
         Args:
             initializer (tf.keras.initializers.Initializer, optional): Initializer used to initialize the scalar weight.
-                Defaults to tf.keras.initializers.Ones().
+                Defaults to None..
         """
         super(ScaleLayer, self).__init__()
-        self.initializer = initializer
+        self.initializer = tf.keras.initializers.Ones() if initializer is None else initializer
         self.scale = self.add_weight(shape=[1], initializer=initializer, constraint=NonNegative(), trainable=True)
 
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
@@ -1147,8 +1146,8 @@ def __init__(
         self,
         use_spec_norm: bool = False,
         intermediate_channel: Optional[bool] = None,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
         """Initializes the `ScalarGatedSelfAttention` layer.
 
@@ -1158,15 +1157,15 @@ def __init__(
             intermediate_channel (Optional[bool], optional): Intermediate channels for the self attention mechanism.
                 Defaults to None.
             kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the kernel weights.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
             gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-                Defaults to HeAlphaUniform().
+                Defaults to None.
         """
         super(ScalarGatedSelfAttention, self).__init__()
         self.use_spec_norm = use_spec_norm
         self.intermediateChannel = intermediate_channel
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
         """Build layer depending on the `input_shape` (output shape of the previous layer).
diff --git a/DeepSaki/layers/padding.py b/DeepSaki/layers/padding.py
index 0f9e300..ac246d4 100644
--- a/DeepSaki/layers/padding.py
+++ b/DeepSaki/layers/padding.py
@@ -7,8 +7,8 @@
 import tensorflow as tf
 
 class ReflectionPadding2D(tf.keras.layers.Layer):
-    """Reflection Padding layer with support for TPU.
-    """
+    """Reflection Padding layer with support for TPU."""
+
     def __init__(self, padding: Tuple[int, int] = (1, 1), **kwargs: Any) -> None:
         """Initialize the `ReflectionPadding2D` layer.
 
diff --git a/DeepSaki/layers/pooling.py b/DeepSaki/layers/pooling.py
index 9948b40..1ed378b 100644
--- a/DeepSaki/layers/pooling.py
+++ b/DeepSaki/layers/pooling.py
@@ -16,6 +16,7 @@ class FrequencyFilter(Enum):
         HIGH_PASS (int): Indicates that high frequency components shall be kept and low frequency components shall be
             filtered.
     """
+
     LOW_PASS = 1
     HIGH_PASS = 2
 
diff --git a/DeepSaki/layers/sub_model_composites.py b/DeepSaki/layers/sub_model_composites.py
index e23518e..918ce89 100644
--- a/DeepSaki/layers/sub_model_composites.py
+++ b/DeepSaki/layers/sub_model_composites.py
@@ -19,31 +19,13 @@
 from DeepSaki.layers.layer_helper import dropout_func
 
 class Encoder(tf.keras.layers.Layer):
-    """
-    Encoder sub-model combines convolutional blocks with down sample blocks. The spatial width is halfed with every level while the channel depth is doubled.
-    args:
-      - number_of_levels (optional, default:3): number of conv2D -> Downsampling pairs
-      - filters (optional, default:64): defines the number of filters to which the input is exposed.
-      - kernels: size of the convolutions kernels
-      - limit_filters (optional, default:1024): limits the number of filters, which is doubled with every downsampling block
-      - use_residual_Conv2DBlock (optional, default: False): ads a residual connection in parallel to the Conv2DBlock
-      - downsampling(optional, default: "conv_stride_2"): describes the downsampling method used
-      - split_kernels (optional, default: False): to decrease the number of parameters, a convolution with the kernel_size (kernel,kernel) can be splitted into two consecutive convolutions with the kernel_size (kernel,1) and (1,kernel) respectivly
-      - number_of_convs (optional, default: 1): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
-      - activation (optional, default: "leaky_relu"): string literal or tensorflow activation function object to obtain activation function
-      - first_kernel (optional, default: 5): The first convolution can have a different kernel size, to e.g. increase the perceptive field, while the channel depth is still low.
-      - use_ResidualBlock (optional, default: False): Whether or not to use the ResidualBlock instead of the Conv2DBlock
-      - residual_cardinality (optional, default: 1): cardinality for the ResidualBlock
-      - channel_list (optional, default:None): alternativly to number_of_layers and filters, a list with the disired filters for each level can be provided. e.g. channel_list = [64, 128, 256] results in a 3-level Encoder with 64, 128, 256 filters for level 1, 2 and 3 respectivly.
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - dropout_rate (optional, default: 0): probability of the dropout layer. If the preceeding layer has more than one channel, spatial dropout is applied, otherwise standard dropout
-      - use_self_attention (optional, default: False): Determines whether to apply self-attention after the encoder before branching.
-      - omit_skips (optional, default: 0): defines how many layers should not output a skip connection output. Requires output_skips to be True. E.g. if omit_skips = 2, the first two levels do not output a skip connection, it starts at level 3.
-      - padding (optional, default: "none"): padding type. Options are "none", "zero" or "reflection"
-      - output_skips (optional, default: False): Whether or not to output skip connections at each level
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
+    """Combines conv blocks with down sample blocks.
+
+    The spatial width is halved with every level while the channel depth is doubled.
+    Can be combined with `dsk.layers.Decoder` and `dsk.layers.Bottleneck` to form an auto encoder model.
+
+    Tipp:
+        Checkout the dsk.models api to find models using this layer.
     """
 
     def __init__(
@@ -68,9 +50,52 @@ def __init__(
         omit_skips: int = 0,
         padding: PaddingType = PaddingType.ZERO,
         output_skips: bool = False,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
+        """Initializes the `Encoder` layer.
+
+        Args:
+            number_of_levels (int, optional): Number of downsampling levels of the model. Defaults to 3.
+            filters (int, optional): Number of filters for the initial encoder block. Defaults to 64.
+            limit_filters (int, optional): Limits the number of filters, which is doubled with every downsampling block.
+                Defaults to 1024.
+            use_residual_Conv2DBlock (bool, optional): Adds a residual connection in parallel to the `Conv2DBlock`.
+                Defaults to False.
+            downsampling (str, optional): Describes the downsampling method. Defaults to "conv_stride_2".
+            kernels (int, optional): Size of the convolutions kernels. Defaults to 3.
+            split_kernels (bool, optional): To decrease the number of parameters, a convolution with the kernel_size
+                `(kernel,kernel)` can be splitted into two consecutive convolutions with the kernel_size `(kernel,1)` and
+                `(1,kernel)` respectivly. Defaults to False.
+            number_of_convs (int, optional): Number of consecutive convolutional building blocks, i.e. `Conv2DBlock`.
+                Defaults to 2.
+            activation (str, optional): String literal or tensorflow activation function object to obtain activation
+                function. Defaults to "leaky_relu".
+            first_kernel (Optional[int], optional): The first convolution can have a different kernel size, to e.g.
+                increase the perceptive field, while the channel depth is still low. Defaults to None.
+            use_ResidualBlock (bool, optional): Whether or not to use the `ResidualBlock` instead of the
+                `Conv2DBlock`. Defaults to False.
+            residual_cardinality (int, optional): Cardinality for the `ResidualBlock`. Defaults to 1.
+            channel_list (Optional[List[int]], optional): alternativly to number_of_layers and filters, a list with the
+                disired filters for each level can be provided. e.g. channel_list = [64, 128, 256] results in a 3-level
+                Encoder with 64, 128, 256 filters for level 1, 2 and 3 respectivly. Defaults to None.
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers.
+                Defaults to False.
+            use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
+            dropout_rate (float, optional): Probability of the dropout layer. If the preceeding layer has more than one
+                channel, spatial dropout is applied, otherwise standard dropout. Defaults to 0.0.
+            use_self_attention (bool, optional): Determines whether to apply self-attention in the encoder. Defaults to False.
+            omit_skips (int, optional): Defines how many layers should not output a skip connection output. Requires
+                `output_skips` to be True. E.g. if `omit_skips = 2`, the first two levels do not output a skip connection,
+                it starts at level 3. Defaults to 0.
+            padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
+            output_skips (bool, optional): If true, ski connections are output. Could be used to attach an encoder to a
+                decoder model. Defaults to False.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to None.
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to None.
+        """
         super(Encoder, self).__init__()
         self.number_of_levels = number_of_levels
         self.filters = filters
@@ -92,10 +117,15 @@ def __init__(
         self.padding = padding
         self.output_skips = output_skips
         self.use_bias = use_bias
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(Encoder, self).build(input_shape)
 
         if self.channel_list is None:
@@ -233,31 +263,16 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
-# Testcode
-# layer = Encoder( number_of_levels = 5, filters = 64, limit_filters = 512, use_self_attention = True,use_residual_Conv2DBlock = True, downsampling="max_pooling", kernels=3, split_kernels = True,  number_of_convs = 2,activation = "leaky_relu", first_kernel=3,use_ResidualBlock = True,use_spec_norm=True, omit_skips=2)
-# print(layer.get_config())
-# dsk.layers.plot_layer(layer,input_shape=(256,256,4))
+class Bottleneck(tf.keras.layers.Layer):
+    """Bottlenecks are sub-model blocks in auto-encoder-like models such as UNet or ResNet.
 
+    It is composed of multiple convolution blocks which might have residuals.
+
+    Can be combined with `dsk.layers.Encoder` and `dsk.layers.Decoder` to form an auto encoder model.
+
+    Tipp:
+        Checkout the dsk.models api to find models using this layer.
 
-class Bottleneck(tf.keras.layers.Layer):
-    """
-    Bottlenecks are sub-model blocks in auto-encoder-like models such as UNet or ResNet. It is composed of multiple convolution blocks which might have residuals
-    args:
-      - n_bottleneck_blocks (optional, default: 3): Number of consecutive convolution blocks
-      - kernels: size of the convolutions kernels
-      - split_kernels (optional, default: False): to decrease the number of parameters, a convolution with the kernel_size (kernel,kernel) can be splitted into two consecutive convolutions with the kernel_size (kernel,1) and (1,kernel) respectivly
-      - number_of_convs (optional, default: 2): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
-      - use_residual_Conv2DBlock (optional, default: True): ads a residual connection in parallel to the Conv2DBlock
-      - use_ResidualBlock (optional, default: False): Whether or not to use the ResidualBlock instead of the Conv2DBlock
-      - activation (optional, default: "leaky_relu"): string literal or tensorflow activation function object to obtain activation function
-      - dropout_rate (optional, default: 0): probability of the dropout layer. If the preceeding layer has more than one channel, spatial dropout is applied, otherwise standard dropout
-      - channel_list (optional, default:None): alternativly to number_of_layers and filters, a list with the disired filters for each block can be provided. e.g. channel_list = [64, 128, 256] results in a 3-staged Bottleneck with 64, 128, 256 filters for stage 1, 2 and 3 respectivly.
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - residual_cardinality (optional, default: 1): cardinality for the ResidualBlock
-      - padding (optional, default: "none"): padding type. Options are "none", "zero" or "reflection"
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
     """
 
     def __init__(
@@ -275,9 +290,40 @@ def __init__(
         use_bias: bool = True,
         residual_cardinality: int = 1,
         padding: PaddingType = PaddingType.ZERO,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
+        """Initializes the `Bottleneck` layer.
+
+        Args:
+            n_bottleneck_blocks (int, optional): Number of consecutive blocks. Defaults to 3.
+            kernels (int, optional): Size of the convolutions kernels. Defaults to 3.
+            split_kernels (bool, optional): To decrease the number of parameters, a convolution with the kernel_size
+                `(kernel,kernel)` can be splitted into two consecutive convolutions with the kernel_size `(kernel,1)`
+                and `(1,kernel)` respectivly. Defaults to False.
+            number_of_convs (int, optional): : Number of consecutive conv layers within a basic building block.
+                Defaults to 2.
+            use_residual_Conv2DBlock (bool, optional): Adds a residual connection in parallel to the `Conv2DBlock`.
+                Defaults to True.
+            use_ResidualBlock (bool, optional): Whether or not to use the `ResidualBlock` instead of the
+                `Conv2DBlock`. Defaults to False.
+            activation (str, optional): String literal or tensorflow activation function object to obtain activation
+                function. Defaults to "leaky_relu".
+            dropout_rate (float, optional): Probability of the dropout layer. If the preceeding layer has more than one
+                channel, spatial dropout is applied, otherwise standard dropout. Defaults to 0.2.
+            channel_list (Optional[List[int]], optional): Alternativly to number_of_layers and filters, a list with the
+                disired filters for each level can be provided. e.g. channel_list = [64, 128, 256] results in a 3-level
+                Decoder with 64, 128, 256 filters for level 1, 2 and 3 respectivly. Defaults to None.
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers. Defaults
+                to False.
+            use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
+            residual_cardinality (int, optional): Cardinality for the `ResidualBlock`. Defaults to 1.
+            padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to None.
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to None.
+        """
         super(Bottleneck, self).__init__()
         self.use_ResidualBlock = use_ResidualBlock
         self.n_bottleneck_blocks = n_bottleneck_blocks
@@ -296,6 +342,11 @@ def __init__(
         self.gamma_initializer = gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(Bottleneck, self).build(input_shape)
 
         if self.channel_list is None:
@@ -380,36 +431,14 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
-# Testcode
-# layer = Bottleneck(True, 3, False, 3,False,1, "leaky_relu" , dropout_rate = 0.2, channel_list = None)
-# print(layer.get_config())
-# dsk.layers.plot_layer(layer,input_shape=(256,256,64))
+class Decoder(tf.keras.layers.Layer):
+    """Combines conv blocks with up sample blocks.
 
+    The spatial width is doubled with every level while the channel depth is halfed.
+    Can be combined with `dsk.layers.Encoder` and `dsk.layers.Bottleneck` to form an auto encoder model.
 
-class Decoder(tf.keras.layers.Layer):
-    """
-    Decoder sub-model combines convolutional blocks with up sample blocks. The spatial width is double with every level while the channel depth is halfed.
-    args:
-      - number_of_levels (optional, default:3): number of conv2D -> Upsampling pairs
-      - upsampling(optional, default: "2D_upsample_and_conv"): describes the upsampling method used
-      - filters (optional, default:64): defines the number of filters to which the input is exposed.
-      - limit_filters (optional, default:1024): limits the number of filters
-      - use_residual_Conv2DBlock (optional, default: False): ads a residual connection in parallel to the Conv2DBlock
-      - kernels: size of the convolutions kernels
-      - split_kernels (optional, default: False): to decrease the number of parameters, a convolution with the kernel_size (kernel,kernel) can be splitted into two consecutive convolutions with the kernel_size (kernel,1) and (1,kernel) respectivly
-      - number_of_convs (optional, default: 1): number of consecutive convolutional building blocks, i.e. Conv2DBlock.
-      - activation (optional, default: "leaky_relu"): string literal or tensorflow activation function object to obtain activation function
-      - dropout_rate (optional, default: 0): probability of the dropout layer. If the preceeding layer has more than one channel, spatial dropout is applied, otherwise standard dropout. In the decoder only applied to the first half of levels.
-      - use_ResidualBlock (optional, default: False): Whether or not to use the ResidualBlock instead of the Conv2DBlock
-      - residual_cardinality (optional, default: 1): cardinality for the ResidualBlock
-      - channel_list (optional, default:None): alternativly to number_of_layers and filters, a list with the disired filters for each level can be provided. e.g. channel_list = [64, 128, 256] results in a 3-level Decoder with 64, 128, 256 filters for level 1, 2 and 3 respectivly.
-      - use_spec_norm (optional, default: False): applies spectral normalization to convolutional and dense layers
-      - use_bias (optional, default: True): determines whether convolutions and dense layers include a bias or not
-      - use_self_attention (optional, default: False): Determines whether to apply self-attention after the encoder before branching.
-      - enable_skip_connections_input (optional, default: False): Whether or not to input skip connections at each level
-      - padding (optional, default: "none"): padding type. Options are "none", "zero" or "reflection"
-      - kernel_initializer (optional, default: HeAlphaUniform()): Initialization of the convolutions kernels.
-      - gamma_initializer (optional, default: HeAlphaUniform()): Initialization of the normalization layers.
+    Tipp:
+        Checkout the dsk.models api to find models using this layer.
     """
 
     def __init__(
@@ -432,9 +461,47 @@ def __init__(
         use_self_attention: bool = False,
         enable_skip_connections_input: bool = False,
         padding: PaddingType = PaddingType.ZERO,
-        kernel_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
-        gamma_initializer: tf.keras.initializers.Initializer = HeAlphaUniform(),
+        kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+        gamma_initializer: Optional[tf.keras.initializers.Initializer] = None,
     ) -> None:
+        """Initializes the `Decoder` layer.
+
+        Args:
+            number_of_levels (int, optional): Number levels in the decoder. Effectivly the number of convolution and
+                upsample pairs. Defaults to 3.
+            upsampling (str, optional): Describes the upsampling method used. Defaults to "2D_upsample_and_conv".
+            filters (int, optional): Base size of filters the is doubled with every level of the decoder.
+                Defaults to 64.
+            Limits the number of filters, which is doubled with every downsampling block.
+            Defaults to 512.
+            use_residual_Conv2DBlock (bool, optional):Adds a residual connection in parallel to the `Conv2DBlock`.
+                Defaults to False.
+            kernels (int, optional): Size of the convolutions kernels. Defaults to 3.
+            split_kernels (bool, optional): To decrease the number of parameters, a convolution with the kernel_size
+                `(kernel,kernel)` can be splitted into two consecutive convolutions with the kernel_size `(kernel,1)` and
+                `(1,kernel)` respectivly. Defaults to False.
+            number_of_convs (int, optional): Number of consecutive convolutional building blocks, i.e. `Conv2DBlock`.
+                Defaults to 2.
+            activation (str, optional): String literal or tensorflow activation function object to obtain activation
+                function. Defaults to "leaky_relu".
+            dropout_rate (float, optional): Probability of the dropout layer. If the preceeding layer has more than one
+                channel, spatial dropout is applied, otherwise standard dropout. Defaults to 0.2.
+            use_ResidualBlock (bool, optional): Whether or not to use the `ResidualBlock` instead of the
+                `Conv2DBlock`. Defaults to False.
+            residual_cardinality (int, optional): Cardinality for the `ResidualBlock`. Defaults to 1.
+            channel_list (Optional[List[int]], optional): Alternativly to number_of_layers and filters, a list with the
+                disired filters for each level can be provided. e.g. channel_list = [64, 128, 256] results in a 3-level
+                Decoder with 64, 128, 256 filters for level 1, 2 and 3 respectivly. Defaults to None.
+            use_spec_norm (bool, optional): Applies spectral normalization to convolutional and dense layers. Defaults to False.
+            use_bias (bool, optional): Whether convolutions and dense layers include a bias or not. Defaults to True.
+            use_self_attention (bool, optional): Determines whether to apply self-attention in the decoder. Defaults to False.
+            enable_skip_connections_input (bool, optional): Whether or not to input skip connections at each level. Defaults to False.
+            padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
+            kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
+                Defaults to None.
+            gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
+                Defaults to None.
+        """
         super(Decoder, self).__init__()
         self.number_of_levels = number_of_levels
         self.filters = filters
@@ -454,10 +521,16 @@ def __init__(
         self.enable_skip_connections_input = enable_skip_connections_input
         self.residual_cardinality = residual_cardinality
         self.padding = padding
-        self.kernel_initializer = kernel_initializer
-        self.gamma_initializer = gamma_initializer
+
+        self.kernel_initializer = HeAlphaUniform() if kernel_initializer is None else kernel_initializer
+        self.gamma_initializer = HeAlphaUniform() if gamma_initializer is None else gamma_initializer
 
     def build(self, input_shape: tf.TensorShape) -> None:
+        """Build layer depending on the `input_shape` (output shape of the previous layer).
+
+        Args:
+            input_shape (tf.TensorShape): Shape of the input tensor to this layer.
+        """
         super(Decoder, self).build(input_shape)
 
         if self.channel_list is None:
@@ -591,9 +664,3 @@ def get_config(self) -> Dict[str, Any]:
             }
         )
         return config
-
-
-# Testcode
-# layer = Decoder( number_of_levels = 5, filters = 64, limit_filters = 2048, use_self_attention = True,use_residual_Conv2DBlock = False, upsampling="depth_to_space", kernels=3, split_kernels = False,  number_of_convs = 2,activation = "leaky_relu",use_ResidualBlock = True,use_spec_norm=False, dropout_rate = 0.2)
-# print(layer.get_config())
-# dsk.layers.plot_layer(layer,input_shape=(256,256,4))
diff --git a/DeepSaki/models/autoencoders.py b/DeepSaki/models/autoencoders.py
index 56c42f5..7c510a5 100644
--- a/DeepSaki/models/autoencoders.py
+++ b/DeepSaki/models/autoencoders.py
@@ -4,7 +4,7 @@
 from DeepSaki.layers.sub_model_composites import Encoder, Decoder,Bottleneck
 from DeepSaki.layers.layer_composites import Conv2DBlock, DenseBlock, ScalarGatedSelfAttention
 
-from typing import Tuple
+from typing import Tuple, Optional
 
 class UNet(tf.keras.Model):
   """U-Net based autoencoder model with skip conections between encoder and decoder. Input_shape = Output_shape.
@@ -64,8 +64,8 @@ def __init__(self,
             omit_skips:int = 0,
             fully_connected:str = "MLP",
             padding:PaddingType=PaddingType.ZERO,
-            kernel_initializer:tf.keras.initializers.Initializer = HeAlphaUniform(),
-            gamma_initializer:tf.keras.initializers.Initializer =  HeAlphaUniform()
+            kernel_initializer: Optional[tf.keras.initializers.Initializer] = None,
+            gamma_initializer: Optional[tf.keras.initializers.Initializer] =  None
             )->None:
     """Initialize the `UNet` object.
 
@@ -109,9 +109,9 @@ def __init__(self,
             the same result, but linear layers are faster. Option: "MLP" or "1x1_conv". Defaults to "MLP".
         padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
         kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
         gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
     """
     super(UNet, self).__init__()
 
@@ -194,7 +194,7 @@ def __init__(self,
             use_self_attention:bool= False,
             fully_connected:str = "MLP",
             padding:PaddingType=PaddingType.ZERO,
-            kernel_initializer:tf.keras.initializers.Initializer = HeAlphaUniform(),
+            kernel_initializer:Optional[tf.keras.initializers.Initializer] = None,
             gamma_initializer:tf.keras.initializers.Initializer =  HeAlphaUniform()
             ):
     """Initialize the `ResNet` object.
@@ -236,9 +236,9 @@ def __init__(self,
             the same result, but linear layers are faster. Option: "MLP" or "1x1_conv". Defaults to "MLP".
         padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
         kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
         gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
     """
     super(ResNet, self).__init__()
 
diff --git a/DeepSaki/models/discriminators.py b/DeepSaki/models/discriminators.py
index e39ce9e..60e403d 100644
--- a/DeepSaki/models/discriminators.py
+++ b/DeepSaki/models/discriminators.py
@@ -57,7 +57,7 @@ def __init__(self,
             padding:PaddingType=PaddingType.NONE,
             fully_connected:str = "MLP",
             use_self_attention:bool = False,
-            kernel_initializer:tf.keras.initializers.Initializer = HeAlphaUniform(),
+            kernel_initializer:Optional[tf.keras.initializers.Initializer] = None,
             gamma_initializer:tf.keras.initializers.Initializer =  HeAlphaUniform()
             ):
     """Initialize the `LayoutContentDiscriminator` object.
@@ -86,9 +86,9 @@ def __init__(self,
         use_self_attention (bool, optional): Determines whether to apply self-attention in the decoder. Defaults to
             False.
         kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
         gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
 
     Raises:
         ValueError: If provided parameter for `fully_connected` is not supported.
@@ -190,7 +190,7 @@ def __init__(self,
             use_bias:bool = True,
             use_self_attention:bool=False,
             padding:PaddingType=PaddingType.NONE,
-            kernel_initializer:tf.keras.initializers.Initializer = HeAlphaUniform(),
+            kernel_initializer:Optional[tf.keras.initializers.Initializer] = None,
             gamma_initializer:tf.keras.initializers.Initializer =  HeAlphaUniform()
             ):
     """Initialize the `PatchDiscriminator` object.
@@ -219,9 +219,9 @@ def __init__(self,
             False.
         padding (PaddingType, optional): Padding type. Defaults to PaddingType.NONE.
         kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
         gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
     """
     super(PatchDiscriminator, self).__init__()
 
@@ -307,7 +307,7 @@ def __init__(self,
             use_bias:bool= True,
             fully_connected:str = "MLP",
             padding:PaddingType=PaddingType.ZERO,
-            kernel_initializer:tf.keras.initializers.Initializer = HeAlphaUniform(),
+            kernel_initializer:Optional[tf.keras.initializers.Initializer] = None,
             gamma_initializer:tf.keras.initializers.Initializer =  HeAlphaUniform()
             ):
     """Initialize the `UNetDiscriminator` object.
@@ -346,9 +346,9 @@ def __init__(self,
             the same result, but linear layers are faster. Option: "MLP" or "1x1_conv". Defaults to "MLP".
         padding (PaddingType, optional): Padding type. Defaults to PaddingType.ZERO.
         kernel_initializer (tf.keras.initializers.Initializer, optional): Initialization of the convolutions kernels.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
         gamma_initializer (tf.keras.initializers.Initializer, optional): Initialization of the normalization layers.
-            Defaults to HeAlphaUniform().
+            Defaults to None.
     """
     super(UNetDiscriminator, self).__init__()
     ch = filters

From 628280874102546719c6eb04bfb0cc45283d70bb Mon Sep 17 00:00:00 2001
From: sascha-kirch <susch130993@googlemail.com>
Date: Fri, 20 Oct 2023 17:02:45 +0200
Subject: [PATCH 6/9] update ruff rules

---
 pyproject.toml | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index af18cb8..ab2e9b0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,7 +15,7 @@ authors = [
 maintainers = [
     {name = "Sascha Kirch", email = "susch130993@googlemail.com"},
 ]
-description = "Custom deep learning code for TensorFlow."
+description = "Collection of useful deep learning custom implementations to extend TensorFlow."
 keywords = ["deeplearning", "machinelearning", "tensorflow","TPU"]
 license = {file = "LICENSE"}
 requires-python = ">=3.11"
@@ -40,10 +40,10 @@ dependencies = [
 dynamic = ["version","readme"]
 
 # https://packaging.python.org/en/latest/specifications/declaring-project-metadata/#urls
-#[project.urls]
+[project.urls]
 #Homepage = "https://example.com"
-#Documentation = "https://readthedocs.org"
-#Repository = "https://github.com/me/spam.git"
+Documentation = "https://sascha-kirch.github.io/DeepSaki"
+Repository = "https://github.com/sascha-kirch/DeepSaki"
 #Changelog = "https://github.com/me/spam/blob/master/CHANGELOG.md"
 
 [project.optional-dependencies]
@@ -176,9 +176,9 @@ exclude = ["examples"]
 [tool.ruff]
 target-version = "py311"
 select = ["F", "E", "I", "D", "PL","N"]
-extend-select = ["W", "C90", "YTT", "ANN", "ASYNC", "A", "C4", "RET", "SIM", "PLC", "PLE", "PLR", "PLW", "FLY", "PERF", "RUF","ARG","ERA","TRY","NPY","FURB", "B","COM","PYI","PT","Q","RET", "ICN"]
+extend-select = ["W", "C90", "YTT", "ANN", "ASYNC", "A", "C4", "RET", "SIM", "PLC", "PLE", "PLR", "PLW", "FLY", "PERF", "RUF","ARG","ERA","TRY","NPY","FURB", "B","COM","PYI","PT","Q","RET", "ICN","FIX","SLF","PIE","FBT","BLE","S"]
 ignore = ["E402","E501", "B905", "SIM300", "PD901","ANN101","COM812"]
-fixable = ["W", "C90", "YTT", "ANN", "ASYNC", "A", "C4", "RET", "SIM", "PLC", "PLE", "PLR", "PLW", "FLY", "PERF", "RUF", "F401", "D205", "D403", "D200", "D202", "D209","I001","E711","Q","COM819","Q002"]
+fixable = ["W", "C90", "YTT", "ANN", "ASYNC", "A", "C4", "RET", "PLC", "PLE", "PLR", "PLW", "FLY", "PERF", "RUF", "F401", "D205", "D403", "D200", "D202", "D209","I001","E711","Q","COM819","Q002"]
 unfixable = ["SIM"]
 line-length = 120
 extend-exclude = ["tests", "test","__init__.py"]
@@ -217,7 +217,7 @@ max-statements = 100
 max-args = 15
 allow-magic-value-types = ["str", "bytes", "complex", "float", "int", "tuple"]
 [tool.ruff.pep8-naming]
-ignore-names = ["iFFT2D","rFFT2D","use_residual_Conv2DBlock","use_ResidualIdentityBlock","rFFT2DFilter"]
+ignore-names = ["iFFT2D","iFFT3D","rFFT2D","use_residual_Conv2DBlock","use_ResidualBlock","rFFT2DFilter"]
 
 # Static type checker
 [tool.pyright]

From c78450f1ec3181a0018885b2be8fda3bbc66f7a4 Mon Sep 17 00:00:00 2001
From: sascha-kirch <susch130993@googlemail.com>
Date: Fri, 20 Oct 2023 21:08:58 +0200
Subject: [PATCH 7/9] fix test case after change of activation func

---
 .../activations_test/complex_valued_activations_test.py       | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/DeepSaki_test/activations_test/complex_valued_activations_test.py b/tests/DeepSaki_test/activations_test/complex_valued_activations_test.py
index 50d49da..c3d1e41 100644
--- a/tests/DeepSaki_test/activations_test/complex_valued_activations_test.py
+++ b/tests/DeepSaki_test/activations_test/complex_valued_activations_test.py
@@ -56,7 +56,7 @@ def test_call_has_expected_output(self, activation, input, expected):
         ],
     )
     def test_dtype(self, input):
-        activation = ComplexActivation()
+        activation = ComplexActivation(tf.keras.layers.ReLU())
         output = activation(input)
         assert output.dtype in (tf.dtypes.complex64, tf.dtypes.complex128)
 
@@ -78,6 +78,6 @@ def test_dtype(self, input):
         ],
     )
     def test_input_shape_equals_output_shape(self, input):
-        activation = ComplexActivation()
+        activation = ComplexActivation(tf.keras.layers.ReLU())
         output = activation(input)
         assert output.shape == input.shape

From ff9704bbf7c4978c237efb68810909b6512774fd Mon Sep 17 00:00:00 2001
From: sascha-kirch <susch130993@googlemail.com>
Date: Fri, 20 Oct 2023 21:11:25 +0200
Subject: [PATCH 8/9] Misc - rename loss->losses - rename optimizer->optimizers
 - add docstrings - fix some warnings from linter

---
 DeepSaki/__init__.py                          |   8 +-
 DeepSaki/augmentations/grid_cutting.py        |   7 +-
 DeepSaki/layers/layer_composites.py           |  90 ++++++------
 DeepSaki/layers/layer_helper.py               |  13 +-
 DeepSaki/layers/sub_model_composites.py       |  21 ++-
 DeepSaki/loss/__init__.py                     |   2 -
 DeepSaki/losses/__init__.py                   |   2 +
 .../{loss => losses}/image_based_losses.py    |   0
 DeepSaki/models/autoencoders.py               |   2 +-
 DeepSaki/models/discriminators.py             |   3 +-
 DeepSaki/optimizer/__init__.py                |   2 -
 DeepSaki/{optimizer => optimizers}/README.md  |   0
 DeepSaki/optimizers/__init__.py               |   2 +
 DeepSaki/{optimizer => optimizers}/swats.py   | 131 +++++++++---------
 DeepSaki/utils/environment.py                 |   9 +-
 README.md                                     |   8 +-
 mkdocs.yml                                    |   2 +-
 pyproject.toml                                |   2 +-
 18 files changed, 164 insertions(+), 140 deletions(-)
 delete mode 100644 DeepSaki/loss/__init__.py
 create mode 100644 DeepSaki/losses/__init__.py
 rename DeepSaki/{loss => losses}/image_based_losses.py (100%)
 delete mode 100644 DeepSaki/optimizer/__init__.py
 rename DeepSaki/{optimizer => optimizers}/README.md (100%)
 create mode 100644 DeepSaki/optimizers/__init__.py
 rename DeepSaki/{optimizer => optimizers}/swats.py (82%)

diff --git a/DeepSaki/__init__.py b/DeepSaki/__init__.py
index 0c0f049..8cda221 100644
--- a/DeepSaki/__init__.py
+++ b/DeepSaki/__init__.py
@@ -3,9 +3,9 @@
 from . import constraints
 from . import initializers
 from . import layers
-from . import loss
+from . import losses
 from . import models
-from . import optimizer
+from . import optimizers
 from . import utils
 
 __version__ = "1.0.0"
@@ -18,8 +18,8 @@
     "constraints",
     "initializers",
     "layers",
-    "loss",
+    "losses",
     "models",
-    "optimizer",
+    "optimizers",
     "utils",
 ]
diff --git a/DeepSaki/augmentations/grid_cutting.py b/DeepSaki/augmentations/grid_cutting.py
index 0a2e4bd..3214206 100644
--- a/DeepSaki/augmentations/grid_cutting.py
+++ b/DeepSaki/augmentations/grid_cutting.py
@@ -18,11 +18,12 @@ def _random_boundingbox(height: int, width: int) -> Tuple[int, int, int, int]:
         x value bottom-right corner.
         y value bottom-right corner.
     """
-    r = np.sqrt(1.0 - np.random.beta(1, 1))
+    random_generator = np.random.default_rng()
+    r = np.sqrt(1.0 - random_generator.beta(1, 1))
     w = np.floor(width * r)
     h = np.floor(height * r)
-    x = np.random.randint(width)
-    y = np.random.randint(height)
+    x = random_generator.integers(width)
+    y = random_generator.integers(height)
 
     x1 = int(np.clip(x - w // 2, 0, width))
     y1 = int(np.clip(y - h // 2, 0, height))
diff --git a/DeepSaki/layers/layer_composites.py b/DeepSaki/layers/layer_composites.py
index 85b13c5..9e9a4f6 100644
--- a/DeepSaki/layers/layer_composites.py
+++ b/DeepSaki/layers/layer_composites.py
@@ -689,51 +689,51 @@ def __init__(
         # for each block, add several con
         self.blocks = []
         for _ in range(number_of_blocks):
-            cardinals = []
-            for _ in range(residual_cardinality):
-                cardinals.append(
-                    [
-                        Conv2DBlock(
-                            filters=self.intermediateFilters,
-                            use_residual_Conv2DBlock=False,
-                            kernels=1,
-                            split_kernels=False,
-                            number_of_convs=1,
-                            activation=activation,
-                            use_spec_norm=use_spec_norm,
-                            use_bias=use_bias,
-                            padding=padding,
-                            kernel_initializer=self.kernel_initializer,
-                            gamma_initializer=self.gamma_initializer,
-                        ),
-                        Conv2DBlock(
-                            filters=self.intermediateFilters,
-                            use_residual_Conv2DBlock=False,
-                            kernels=kernels,
-                            split_kernels=False,
-                            number_of_convs=1,
-                            activation=activation,
-                            padding=PaddingType.NONE,
-                            use_spec_norm=use_spec_norm,
-                            use_bias=use_bias,
-                            kernel_initializer=self.kernel_initializer,
-                            gamma_initializer=self.gamma_initializer,
-                        ),
-                        Conv2DBlock(
-                            filters=filters,
-                            use_residual_Conv2DBlock=False,
-                            kernels=1,
-                            split_kernels=False,
-                            number_of_convs=1,
-                            activation=activation,
-                            use_spec_norm=use_spec_norm,
-                            use_bias=use_bias,
-                            padding=padding,
-                            kernel_initializer=self.kernel_initializer,
-                            gamma_initializer=self.gamma_initializer,
-                        ),
-                    ]
-                )
+            cardinals = [
+                [
+                    Conv2DBlock(
+                        filters=self.intermediateFilters,
+                        use_residual_Conv2DBlock=False,
+                        kernels=1,
+                        split_kernels=False,
+                        number_of_convs=1,
+                        activation=activation,
+                        use_spec_norm=use_spec_norm,
+                        use_bias=use_bias,
+                        padding=padding,
+                        kernel_initializer=self.kernel_initializer,
+                        gamma_initializer=self.gamma_initializer,
+                    ),
+                    Conv2DBlock(
+                        filters=self.intermediateFilters,
+                        use_residual_Conv2DBlock=False,
+                        kernels=kernels,
+                        split_kernels=False,
+                        number_of_convs=1,
+                        activation=activation,
+                        padding=PaddingType.NONE,
+                        use_spec_norm=use_spec_norm,
+                        use_bias=use_bias,
+                        kernel_initializer=self.kernel_initializer,
+                        gamma_initializer=self.gamma_initializer,
+                    ),
+                    Conv2DBlock(
+                        filters=filters,
+                        use_residual_Conv2DBlock=False,
+                        kernels=1,
+                        split_kernels=False,
+                        number_of_convs=1,
+                        activation=activation,
+                        use_spec_norm=use_spec_norm,
+                        use_bias=use_bias,
+                        padding=padding,
+                        kernel_initializer=self.kernel_initializer,
+                        gamma_initializer=self.gamma_initializer,
+                    ),
+                ]
+                for _ in range(residual_cardinality)
+            ]
+
             self.blocks.append(cardinals)
 
         self.dropout = dropout_func(filters, dropout_rate)
diff --git a/DeepSaki/layers/layer_helper.py b/DeepSaki/layers/layer_helper.py
index 26dcccc..efcb923 100644
--- a/DeepSaki/layers/layer_helper.py
+++ b/DeepSaki/layers/layer_helper.py
@@ -69,7 +69,8 @@ def get_initializer(initializer: InitializerFunc, seed: Optional[int] = None) ->
         InitializerFunc.HE_ALPHA_UNIFORM: HeAlphaUniform(seed=seed),
     }
 
-    assert initializer in valid_options, f"Undefined initializer provided: {initializer}"
+    if initializer not in valid_options:
+        raise ValueError(f"Undefined initializer provided: {initializer}")
 
     return valid_options.get(initializer)
 
@@ -81,7 +82,7 @@ def pad_func(
 
     Args:
         pad_values (Tuple[int,int], optional): Size of the padding values. Defaults to (1, 1).
-        padding_type (PaddingType, optional): [_description_]. Defaults to PaddingType.ZERO.
+        padding_type (PaddingType, optional): Padding Type. Defaults to PaddingType.ZERO.
 
     Returns:
         Instance of a padding layer object.
@@ -90,9 +91,11 @@ def pad_func(
         PaddingType.REFLECTION: ReflectionPadding2D(pad_values),
         PaddingType.ZERO: tf.keras.layers.ZeroPadding2D(pad_values),
     }
-    assert (
-        padding_type in valid_options
-    ), f"Undefined padding type provided: '{padding_type}'. Valid options are: '{valid_options.keys()}'"
+    if padding_type not in valid_options:
+        raise ValueError(
+            f"Undefined padding type provided: '{padding_type}'. Valid options are: '{valid_options.keys()}'"
+        )
+
     return valid_options.get(padding_type)
 
 
diff --git a/DeepSaki/layers/sub_model_composites.py b/DeepSaki/layers/sub_model_composites.py
index 918ce89..7b9060c 100644
--- a/DeepSaki/layers/sub_model_composites.py
+++ b/DeepSaki/layers/sub_model_composites.py
@@ -206,6 +206,21 @@ def build(self, input_shape: tf.TensorShape) -> None:
                 )
 
     def call(self, inputs: tf.Tensor) -> Union[tf.Tensor, Tuple[tf.Tensor, tf.Tensor]]:
+        """Calls the `Encoder` layer.
+
+        Args:
+            inputs (tf.Tensor): Input tensor of shape (batch, height, width, channel)
+
+        Raises:
+            ValueError: if layer has not been build by calling build() on to layer.
+
+        Returns:
+            If `output_skips=False` only the final output of the Encoder is returned as a tensor of shape
+                (`batch`, `height/2**number_of_levels`, `width/2**number_of_levels`,
+                `max(channel*2**number_of_levels, limit_filters)`. If `output_skips=True` additionally returns a tensor
+                of tensor (one for each level of the encoder) of shape (`batch`, `height/2**level`, `width/2**level`,
+                `max(channel*2**level, limit_filters)`.
+        """
         if not self.built:
             raise ValueError("This model has not yet been built.")
 
@@ -472,8 +487,8 @@ def __init__(
             upsampling (str, optional): Describes the upsampling method used. Defaults to "2D_upsample_and_conv".
             filters (int, optional): Base size of filters the is doubled with every level of the decoder.
                 Defaults to 64.
-            Limits the number of filters, which is doubled with every downsampling block.
-            Defaults to 512.
+            limit_filters (int, optional): Limits the number of filters, which is doubled with every downsampling block.
+                Defaults to 512.
             use_residual_Conv2DBlock (bool, optional):Adds a residual connection in parallel to the `Conv2DBlock`.
                 Defaults to False.
             kernels (int, optional): Size of the convolutions kernels. Defaults to 3.
@@ -542,7 +557,6 @@ def build(self, input_shape: tf.TensorShape) -> None:
 
         self.decoderBlocks = []
         self.upSampleBlocks = []
-        self.dropouts = []
 
         if self.use_self_attention:
             self.SA = ScalarGatedSelfAttention(
@@ -604,7 +618,6 @@ def build(self, input_shape: tf.TensorShape) -> None:
                     UpSampleBlock(
                         kernels=self.kernels,
                         upsampling=self.upsampling,
-                        split_kernels=self.split_kernels,
                         activation=self.activation,
                         use_spec_norm=self.use_spec_norm,
                         use_bias=self.use_bias,
diff --git a/DeepSaki/loss/__init__.py b/DeepSaki/loss/__init__.py
deleted file mode 100644
index 803796f..0000000
--- a/DeepSaki/loss/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from DeepSaki.loss.image_based_losses import PixelDistanceLoss
-from DeepSaki.loss.image_based_losses import StructuralSimilarityLoss
diff --git a/DeepSaki/losses/__init__.py b/DeepSaki/losses/__init__.py
new file mode 100644
index 0000000..7aff341
--- /dev/null
+++ b/DeepSaki/losses/__init__.py
@@ -0,0 +1,2 @@
+from DeepSaki.losses.image_based_losses import PixelDistanceLoss
+from DeepSaki.losses.image_based_losses import StructuralSimilarityLoss
diff --git a/DeepSaki/loss/image_based_losses.py b/DeepSaki/losses/image_based_losses.py
similarity index 100%
rename from DeepSaki/loss/image_based_losses.py
rename to DeepSaki/losses/image_based_losses.py
diff --git a/DeepSaki/models/autoencoders.py b/DeepSaki/models/autoencoders.py
index 7c510a5..b014a10 100644
--- a/DeepSaki/models/autoencoders.py
+++ b/DeepSaki/models/autoencoders.py
@@ -2,7 +2,7 @@
 from DeepSaki.initializers.he_alpha import HeAlphaUniform
 from DeepSaki.layers.layer_helper import PaddingType
 from DeepSaki.layers.sub_model_composites import Encoder, Decoder,Bottleneck
-from DeepSaki.layers.layer_composites import Conv2DBlock, DenseBlock, ScalarGatedSelfAttention
+from DeepSaki.layers.layer_composites import Conv2DBlock, DenseBlock
 
 from typing import Tuple, Optional
 
diff --git a/DeepSaki/models/discriminators.py b/DeepSaki/models/discriminators.py
index 60e403d..245e349 100644
--- a/DeepSaki/models/discriminators.py
+++ b/DeepSaki/models/discriminators.py
@@ -6,7 +6,7 @@
 from DeepSaki.layers.layer_composites import Conv2DBlock, DenseBlock, ScalarGatedSelfAttention
 from DeepSaki.layers.pooling import GlobalSumPooling2D
 
-from typing import Tuple
+from typing import Tuple, Optional
 
 class LayoutContentDiscriminator(tf.keras.Model):
   """Discriminator/critic model with two outputs to enforce disentanglement of layout and content discrimination.
@@ -249,7 +249,6 @@ def call(self, inputs:tf.Tensor)->tf.Tensor:
 class UNetDiscriminator(tf.keras.Model):
   """U-Net based discriminator for pixel-wise real/fake prediction plus additional output for global prediction.
 
-
   Usually used in a GAN framework.
 
   Info:
diff --git a/DeepSaki/optimizer/__init__.py b/DeepSaki/optimizer/__init__.py
deleted file mode 100644
index 52f920b..0000000
--- a/DeepSaki/optimizer/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from DeepSaki.optimizer.swats import SWATS_ADAM
-from DeepSaki.optimizer.swats import SWATS_NADAM
diff --git a/DeepSaki/optimizer/README.md b/DeepSaki/optimizers/README.md
similarity index 100%
rename from DeepSaki/optimizer/README.md
rename to DeepSaki/optimizers/README.md
diff --git a/DeepSaki/optimizers/__init__.py b/DeepSaki/optimizers/__init__.py
new file mode 100644
index 0000000..11a7424
--- /dev/null
+++ b/DeepSaki/optimizers/__init__.py
@@ -0,0 +1,2 @@
+from DeepSaki.optimizers.swats import SwatsAdam
+from DeepSaki.optimizers.swats import SwatsNadam
diff --git a/DeepSaki/optimizer/swats.py b/DeepSaki/optimizers/swats.py
similarity index 82%
rename from DeepSaki/optimizer/swats.py
rename to DeepSaki/optimizers/swats.py
index 0d06543..44587cf 100644
--- a/DeepSaki/optimizer/swats.py
+++ b/DeepSaki/optimizers/swats.py
@@ -1,10 +1,12 @@
-# The SWATS Optimizer is an optimizer that starts with ADAM and switches to SGD after a certain epoch.
-# SWATS Paper: http://arxiv.org/abs/1712.07628
-#
-# This implementation is based on tensorflow's ADAM, SGD and NADAM optimizer
-# TensorFlow ADAM: https://github.com/tensorflow/tensorflow/blob/85c8b2a817f95a3e979ecd1ed95bff1dc1335cff/tensorflow/python/keras/optimizer_v2/adam.py
-# Tensorflow SGD: https://github.com/tensorflow/tensorflow/blob/85c8b2a817f95a3e979ecd1ed95bff1dc1335cff/tensorflow/python/keras/optimizer_v2/gradient_descent.py
-# Tensorflow NADAM: https://github.com/tensorflow/tensorflow/blob/85c8b2a817f95a3e979ecd1ed95bff1dc1335cff/tensorflow/python/keras/optimizer_v2/nadam.py
+"""The SWATS Optimizer is an optimizer that starts with ADAM and switches to SGD after a certain epoch.
+
+SWATS Paper: http://arxiv.org/abs/1712.07628
+
+This implementation is based on tensorflow's ADAM, SGD and NADAM optimizer
+TensorFlow ADAM: https://github.com/tensorflow/tensorflow/blob/85c8b2a817f95a3e979ecd1ed95bff1dc1335cff/tensorflow/python/keras/optimizer_v2/adam.py
+Tensorflow SGD: https://github.com/tensorflow/tensorflow/blob/85c8b2a817f95a3e979ecd1ed95bff1dc1335cff/tensorflow/python/keras/optimizer_v2/gradient_descent.py
+Tensorflow NADAM: https://github.com/tensorflow/tensorflow/blob/85c8b2a817f95a3e979ecd1ed95bff1dc1335cff/tensorflow/python/keras/optimizer_v2/nadam.py
+"""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -25,7 +27,7 @@
 from tensorflow.python.ops import variables as tf_variables
 from tensorflow.python.training import gen_training_ops
 
-class SWATS_ADAM(optimizer_v2.OptimizerV2):
+class SwatsAdam(optimizer_v2.OptimizerV2):
     _HAS_AGGREGATE_GRAD = True
 
     def __init__(
@@ -37,10 +39,10 @@ def __init__(
         amsgrad: bool = False,
         momentum: float = 0.0,
         nesterov: bool = False,
-        name: str = "SWATS_ADAM",
+        name: str = "SwatsAdam",
         **kwargs: Any,
     ) -> None:
-        super(SWATS_ADAM, self).__init__(name, **kwargs)
+        super(SwatsAdam, self).__init__(name, **kwargs)
         self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
         self._set_hyper("decay", self._initial_decay)
         self._set_hyper("beta_1", beta_1)
@@ -73,7 +75,7 @@ def _create_slots(self, var_list: List[tf.Variable]) -> None:
                 self.add_slot(var, "momentum")
 
     def _prepare_local(self, var_device, var_dtype: tf.DType, apply_state) -> None:
-        super(SWATS_ADAM, self)._prepare_local(var_device, var_dtype, apply_state)
+        super(SwatsAdam, self)._prepare_local(var_device, var_dtype, apply_state)
 
         local_step = math_ops.cast(self.iterations + 1, var_dtype)
         beta_1_t = array_ops.identity(self._get_hyper("beta_1", var_dtype))
@@ -104,7 +106,7 @@ def set_weights(self, weights) -> None:
         num_vars = int((len(params) - 1) / 2)
         if len(weights) == 3 * num_vars + 1:
             weights = weights[: len(params)]
-        super(SWATS_ADAM, self).set_weights(weights)
+        super(SwatsAdam, self).set_weights(weights)
 
     def _resource_apply_dense(self, grad, var, apply_state=None):
         var_device, var_dtype = var.device, var.dtype.base_dtype
@@ -124,12 +126,12 @@ def _resource_apply_dense(self, grad, var, apply_state=None):
                     use_locking=self._use_locking,
                     use_nesterov=self.nesterov,
                 )
-            else:
-                return gen_training_ops.ResourceApplyGradientDescent(
-                    var=var.handle, alpha=coefficients["lr_t"], delta=grad, use_locking=self._use_locking
-                )
 
-        elif self.currentOptimizer == "adam":
+            return gen_training_ops.ResourceApplyGradientDescent(
+                var=var.handle, alpha=coefficients["lr_t"], delta=grad, use_locking=self._use_locking
+            )
+
+        if self.currentOptimizer == "adam":
             m = self.get_slot(var, "m")
             v = self.get_slot(var, "v")
             if not self.amsgrad:
@@ -146,24 +148,24 @@ def _resource_apply_dense(self, grad, var, apply_state=None):
                     grad=grad,
                     use_locking=self._use_locking,
                 )
-            else:
-                vhat = self.get_slot(var, "vhat")
-                return gen_training_ops.ResourceApplyAdamWithAmsgrad(
-                    var=var.handle,
-                    m=m.handle,
-                    v=v.handle,
-                    vhat=vhat.handle,
-                    beta1_power=coefficients["beta_1_power"],
-                    beta2_power=coefficients["beta_2_power"],
-                    lr=coefficients["lr_t"],
-                    beta1=coefficients["beta_1_t"],
-                    beta2=coefficients["beta_2_t"],
-                    epsilon=coefficients["epsilon"],
-                    grad=grad,
-                    use_locking=self._use_locking,
-                )
-        else:
-            raise Exception("Optimizer is not Defined. Use adam or sgd.")
+
+            vhat = self.get_slot(var, "vhat")
+            return gen_training_ops.ResourceApplyAdamWithAmsgrad(
+                var=var.handle,
+                m=m.handle,
+                v=v.handle,
+                vhat=vhat.handle,
+                beta1_power=coefficients["beta_1_power"],
+                beta2_power=coefficients["beta_2_power"],
+                lr=coefficients["lr_t"],
+                beta1=coefficients["beta_1_t"],
+                beta2=coefficients["beta_2_t"],
+                epsilon=coefficients["epsilon"],
+                grad=grad,
+                use_locking=self._use_locking,
+            )
+
+        raise Exception("Optimizer is not Defined. Use adam or sgd.")
 
     def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
         var_device, var_dtype = var.device, var.dtype.base_dtype
@@ -184,7 +186,7 @@ def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
                 use_nesterov=self.nesterov,
             )
 
-        elif self.currentOptimizer == "adam":
+        if self.currentOptimizer == "adam":
             m = self.get_slot(var, "m")
             m_scaled_g_values = grad * coefficients["one_minus_beta_1_t"]
             m_t = state_ops.assign(m, m * coefficients["beta_1_t"], use_locking=self._use_locking)
@@ -203,21 +205,21 @@ def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
                     var, coefficients["lr"] * m_t / (v_sqrt + coefficients["epsilon"]), use_locking=self._use_locking
                 )
                 return control_flow_ops.group(*[var_update, m_t, v_t])
-            else:
-                v_hat = self.get_slot(var, "vhat")
-                v_hat_t = math_ops.maximum(v_hat, v_t)
-                with ops.control_dependencies([v_hat_t]):
-                    v_hat_t = state_ops.assign(v_hat, v_hat_t, use_locking=self._use_locking)
-                v_hat_sqrt = math_ops.sqrt(v_hat_t)
-                var_update = state_ops.assign_sub(
-                    var, oefficients["lr"] * m_t / (v_hat_sqrt + coefficients["epsilon"]), use_locking=self._use_locking
-                )
-                return control_flow_ops.group(*[var_update, m_t, v_t, v_hat_t])
-        else:
-            raise Exception("Optimizer is not Defined. Use adam or sgd.")
+
+            v_hat = self.get_slot(var, "vhat")
+            v_hat_t = math_ops.maximum(v_hat, v_t)
+            with ops.control_dependencies([v_hat_t]):
+                v_hat_t = state_ops.assign(v_hat, v_hat_t, use_locking=self._use_locking)
+            v_hat_sqrt = math_ops.sqrt(v_hat_t)
+            var_update = state_ops.assign_sub(
+                var, coefficients["lr"] * m_t / (v_hat_sqrt + coefficients["epsilon"]), use_locking=self._use_locking
+            )
+            return control_flow_ops.group(*[var_update, m_t, v_t, v_hat_t])
+
+        raise Exception("Optimizer is not Defined. Use adam or sgd.")
 
     def get_config(self) -> Dict[str, Any]:
-        config = super(SWATS_ADAM, self).get_config()
+        config = super(SwatsAdam, self).get_config()
         config.update(
             {
                 "learning_rate": self._serialize_hyperparameter("learning_rate"),
@@ -234,7 +236,7 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
 
-class SWATS_NADAM(optimizer_v2.OptimizerV2):
+class SwatsNadam(optimizer_v2.OptimizerV2):
     _HAS_AGGREGATE_GRAD = True
 
     def __init__(
@@ -245,10 +247,10 @@ def __init__(
         epsilon: float = 1e-7,
         momentum: float = 0.0,
         nesterov: bool = False,
-        name: str = "SWATS_NADAM",
+        name: str = "SwatsNadam",
         **kwargs: Any,
     ) -> None:
-        super(SWATS_NADAM, self).__init__(name, **kwargs)
+        super(SwatsNadam, self).__init__(name, **kwargs)
         self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
         self._set_hyper("decay", self._initial_decay)
         self._set_hyper("beta_1", beta_1)
@@ -290,7 +292,7 @@ def _create_slots(self, var_list: List[tf.Variable]) -> None:
                 self.add_slot(var, "momentum")
 
     def _prepare_local(self, var_device, var_dtype: tf.DType, apply_state) -> None:
-        super(SWATS_NADAM, self)._prepare_local(var_device, var_dtype, apply_state)
+        super(SwatsNadam, self)._prepare_local(var_device, var_dtype, apply_state)
 
         # From Nadam Optimizer!
         lr_t = array_ops.identity(self._get_hyper("learning_rate", var_dtype))
@@ -330,7 +332,7 @@ def _prepare_local(self, var_device, var_dtype: tf.DType, apply_state) -> None:
     def _prepare(self, var_list: List[tf.Variable]):
         # Get the value of the momentum cache before starting to apply gradients.
         self._m_cache_read = array_ops.identity(self._m_cache)
-        return super(SWATS_NADAM, self)._prepare(var_list)
+        return super(SwatsNadam, self)._prepare(var_list)
 
     def _resource_apply_dense(self, grad, var, apply_state=None):
         var_device, var_dtype = var.device, var.dtype.base_dtype
@@ -351,12 +353,11 @@ def _resource_apply_dense(self, grad, var, apply_state=None):
                     use_locking=self._use_locking,
                     use_nesterov=self.nesterov,
                 )
-            else:
-                return gen_training_ops.ResourceApplyGradientDescent(
-                    var=var.handle, alpha=coefficients["lr_t"], delta=grad, use_locking=self._use_locking
-                )
+            return gen_training_ops.ResourceApplyGradientDescent(
+                var=var.handle, alpha=coefficients["lr_t"], delta=grad, use_locking=self._use_locking
+            )
         # Nadam Optimizer!
-        elif self.currentOptimizer == "nadam":
+        if self.currentOptimizer == "nadam":
             m = self.get_slot(var, "m")
             v = self.get_slot(var, "v")
             g_prime = grad / coefficients["one_minus_m_schedule_new"]
@@ -369,8 +370,8 @@ def _resource_apply_dense(self, grad, var, apply_state=None):
             m_t_bar = coefficients["one_minus_m_t"] * g_prime + coefficients["m_t_1"] * m_t_prime
             var_t = var - coefficients["lr_t"] * m_t_bar / (math_ops.sqrt(v_t_prime) + coefficients["epsilon"])
             return state_ops.assign(var, var_t, use_locking=self._use_locking).op
-        else:
-            raise Exception("Optimizer is not Defined. Use nadam or sgd.")
+
+        raise Exception("Optimizer is not Defined. Use nadam or sgd.")
 
     def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
         var_device, var_dtype = var.device, var.dtype.base_dtype
@@ -392,7 +393,7 @@ def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
                 use_nesterov=self.nesterov,
             )
         # Nadam Part
-        elif self.currentOptimizer == "nadam":
+        if self.currentOptimizer == "nadam":
             m = self.get_slot(var, "m")
             v = self.get_slot(var, "v")
 
@@ -422,11 +423,11 @@ def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
                 var, indices, coefficients["neg_lr_t"] * m_t_bar / v_prime_sqrt_plus_eps
             )
             return control_flow_ops.group(*[var_update, m_t_bar, v_t])
-        else:
-            raise Exception("Optimizer is not Defined. Use nadam or sgd.")
+
+        raise Exception("Optimizer is not Defined. Use nadam or sgd.")
 
     def get_config(self) -> Dict[str, Any]:
-        config = super(SWATS_NADAM, self).get_config()
+        config = super(SwatsNadam, self).get_config()
         config.update(
             {
                 "learning_rate": self._serialize_hyperparameter("learning_rate"),
diff --git a/DeepSaki/utils/environment.py b/DeepSaki/utils/environment.py
index a8e4343..01fbced 100644
--- a/DeepSaki/utils/environment.py
+++ b/DeepSaki/utils/environment.py
@@ -31,19 +31,20 @@ def detect_accelerator(
         `hw_accelerator_handle`: pointer to the HW accelerator
     """
     hw_accelerator_handle = None
+    tpu = None
+    gpus = []
     try:
         tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
         hw_accelerator_handle = tpu
     except ValueError:
-        tpu = None
         if gpu_memory_groth:
             for gpu in tf.config.experimental.list_physical_devices("GPU"):
-                tf.config.experimental.set_memory_growth(gpu, True)
+                tf.config.experimental.set_memory_growth(gpu, enable=True)
         gpus = tf.config.experimental.list_logical_devices("GPU")
         hw_accelerator_handle = gpus
 
     # Select appropriate distribution strategy
-    if tpu:
+    if tpu is not None:
         runtime_environment = "TPU"
         tf.config.experimental_connect_to_cluster(tpu)
         tf.tpu.experimental.initialize_tpu_system(tpu)
@@ -74,7 +75,7 @@ def detect_accelerator(
 
 def enable_xla_acceleration() -> None:
     """Enable compiler acceleration for linear algebra operations."""
-    tf.config.optimizer.set_jit(True)
+    tf.config.optimizer.set_jit(enabled=True)
     print("Linear algebra acceleration enabled")
 
 
diff --git a/README.md b/README.md
index db76f5c..430fe11 100644
--- a/README.md
+++ b/README.md
@@ -5,10 +5,16 @@ The ML framework used is tensorflow and the entire code is suitable to run Googl
 
 ![Python](https://img.shields.io/badge/python-3.11+-blue)
 ![GitHub](https://img.shields.io/github/license/sascha-kirch/deepsaki)
+
+>Main:<br>
 [![Build](https://github.com/sascha-kirch/DeepSaki/actions/workflows/test.yml/badge.svg?branch=main)](https://github.com/sascha-kirch/DeepSaki/actions/workflows/test.yml)
-[![codecov](https://codecov.io/gh/DeepSaki/branch/main/graph/badge.svg)](https://codecov.io/gh/DeepSaki)
+[![codecov](https://codecov.io/gh/sascha-kirch/DeepSaki/branch/main/graph/badge.svg?token=FD7IE1T9EO)](https://codecov.io/gh/sascha-kirch/DeepSaki)
 [![Documentation](https://img.shields.io/badge/ref-Documentation-blue)](https://sascha-kirch.github.io/DeepSaki/latest/)
 
+>Develop:<br>
+[![Build](https://github.com/sascha-kirch/DeepSaki/actions/workflows/test.yml/badge.svg?branch=develop)](https://github.com/sascha-kirch/DeepSaki/actions/workflows/test.yml)
+[![codecov](https://codecov.io/gh/sascha-kirch/DeepSaki/branch/develop/graph/badge.svg?token=FD7IE1T9EO)](https://codecov.io/gh/sascha-kirch/DeepSaki)
+[![Documentation](https://img.shields.io/badge/ref-Documentation-blue)](https://sascha-kirch.github.io/DeepSaki/develop/)
 
 # Installation
 
diff --git a/mkdocs.yml b/mkdocs.yml
index 0c13263..b82f753 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -17,7 +17,7 @@ theme:
     - navigation.top
     - navigation.footer
     - toc.follow
-    - toc.integrate
+    #- toc.integrate
   name: "material"
   # https://squidfunk.github.io/mkdocs-material/setup/changing-the-colors/
   palette:
diff --git a/pyproject.toml b/pyproject.toml
index ab2e9b0..d51439c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -177,7 +177,7 @@ exclude = ["examples"]
 target-version = "py311"
 select = ["F", "E", "I", "D", "PL","N"]
 extend-select = ["W", "C90", "YTT", "ANN", "ASYNC", "A", "C4", "RET", "SIM", "PLC", "PLE", "PLR", "PLW", "FLY", "PERF", "RUF","ARG","ERA","TRY","NPY","FURB", "B","COM","PYI","PT","Q","RET", "ICN","FIX","SLF","PIE","FBT","BLE","S"]
-ignore = ["E402","E501", "B905", "SIM300", "PD901","ANN101","COM812"]
+ignore = ["E402","E501", "B905", "SIM300", "PD901","ANN101","COM812","FBT001","FBT002"]
 fixable = ["W", "C90", "YTT", "ANN", "ASYNC", "A", "C4", "RET", "PLC", "PLE", "PLR", "PLW", "FLY", "PERF", "RUF", "F401", "D205", "D403", "D200", "D202", "D209","I001","E711","Q","COM819","Q002"]
 unfixable = ["SIM"]
 line-length = 120

From 7cc7505f398c060a0aed9a991a8fbfc35a870556 Mon Sep 17 00:00:00 2001
From: sascha-kirch <susch130993@googlemail.com>
Date: Fri, 20 Oct 2023 21:32:46 +0200
Subject: [PATCH 9/9] add further doc strings

---
 DeepSaki/layers/layer_composites.py     | 14 +++++++++++
 DeepSaki/layers/sub_model_composites.py | 33 ++++++++++++++++++++++---
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/DeepSaki/layers/layer_composites.py b/DeepSaki/layers/layer_composites.py
index 9e9a4f6..624f130 100644
--- a/DeepSaki/layers/layer_composites.py
+++ b/DeepSaki/layers/layer_composites.py
@@ -231,6 +231,20 @@ def build(self, input_shape: tf.TensorShape) -> None:
                 self.residual_conv = tfa.layers.SpectralNormalization(self.residual_conv)
 
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
+        """Calls the `Conv2DBlock` layer.
+
+        Args:
+            inputs (tf.Tensor): Tensor of shape (batch, height, width, channel)
+
+        Raises:
+            ValueError: If layer has not been built by calling build() on to layer.
+
+        Returns:
+            Tensor of shape (batch, `H`, `W`, `filters`). The values for `H`, `W`, `C` depend on the stride as well on
+            the padding. If padding is applied in the stride is (1,1), the output shape matches the input shape. If for
+            example the stride is(2,2) the output shape would be (batch, `height/(2*number_of_convs)`,
+            `width/(2*number_of_convs)`, `filters`).
+        """
         if not self.built:
             raise ValueError("This model has not yet been built.")
         x = inputs
diff --git a/DeepSaki/layers/sub_model_composites.py b/DeepSaki/layers/sub_model_composites.py
index 7b9060c..48e277a 100644
--- a/DeepSaki/layers/sub_model_composites.py
+++ b/DeepSaki/layers/sub_model_composites.py
@@ -212,14 +212,14 @@ def call(self, inputs: tf.Tensor) -> Union[tf.Tensor, Tuple[tf.Tensor, tf.Tensor
             inputs (tf.Tensor): Input tensor of shape (batch, height, width, channel)
 
         Raises:
-            ValueError: if layer has not been build by calling build() on to layer.
+            ValueError: If layer has not been built by calling build() on to layer.
 
         Returns:
             If `output_skips=False` only the final output of the Encoder is returned as a tensor of shape
                 (`batch`, `height/2**number_of_levels`, `width/2**number_of_levels`,
-                `max(channel*2**number_of_levels, limit_filters)`. If `output_skips=True` additionally returns a tensor
+                `min(channel*2**number_of_levels, limit_filters)`. If `output_skips=True` additionally returns a tensor
                 of tensor (one for each level of the encoder) of shape (`batch`, `height/2**level`, `width/2**level`,
-                `max(channel*2**level, limit_filters)`.
+                `min(channel*2**level, limit_filters)`.
         """
         if not self.built:
             raise ValueError("This model has not yet been built.")
@@ -405,6 +405,17 @@ def build(self, input_shape: tf.TensorShape) -> None:
         self.dropout = dropout_func(self.channel_list[-1], self.dropout_rate)
 
     def call(self, inputs: tf.Tensor) -> tf.Tensor:
+        """Calls the `Bottleneck` layer.
+
+        Args:
+            inputs (tf.Tensor): Input tensor of shape (batch, height, width, channel)
+
+        Raises:
+            ValueError: If layer has not been built by calling build() on to layer.
+
+        Returns:
+            Tensor of shape (batch, height, width, channel)
+        """
         if not self.built:
             raise ValueError("This model has not yet been built.")
         x = inputs
@@ -627,7 +638,21 @@ def build(self, input_shape: tf.TensorShape) -> None:
                     )
                 )
 
-    def call(self, inputs: tf.Tensor | Tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
+    def call(self, inputs: Union[tf.Tensor , Tuple[tf.Tensor, tf.Tensor]]) -> tf.Tensor:
+        """Calls the `Decoder` layer.
+
+        Args:
+            inputs (Union[tf.Tensor , Tuple[tf.Tensor, tf.Tensor]]): If `enable_skip_connections_input=False` only
+                inputs a tensor of shape (batch, height, width, min(channel*2**level, limit_filters)). If
+                `enable_skip_connections_input=True`, additonally at every level of the decoder, skip connections from
+                an encoder can be inserted.
+
+        Raises:
+            ValueError: If layer has not been built by calling build() on to layer.
+
+        Returns:
+            tf.Tensor: Tensor of shape (`batch`, `height*2**number_of_levels`, `width*2**number_of_levels`,`filters`).
+        """
         if not self.built:
             raise ValueError("This model has not yet been built.")
         skip_connections = None