layers.py

import tensorflow as tf
from tensorflow import keras

class MonotonicFunction(tf.keras.layers.Layer):
    """
    Pure neural network approach for monotonic water retention curve prediction
    """
    def __init__(self, num_basis=16):
        super(MonotonicFunction, self).__init__()
        self.num_basis = num_basis

    def build(self, input_shape):
        # Initialize basis function weights
        self.basis_weights = self.add_weight(
            shape=(self.num_basis,),
            initializer=tf.keras.initializers.GlorotNormal(),
            trainable=True,
            name='basis_weights'
        )
        
        # Modify output scaling parameters
        # The output scaling mechanism (output_scale and output_shift) constrains predictions 
        # to physically realistic volumetric water content ranges (approximately 0.1-0.6 cm³/cm³).
        # Although these parameters are trainable, the model tends to preserve this realistic range
        # since it aligns with the physical limits of soil water retention (from residual water
        # content to near saturation). If wider ranges are needed, consider adjusting the
        # initializers or adding direct physical constraints instead.
        #
        # Note: (0.6,0.05) or narrower range resulted in tappered model predicted min and max vwc
        self.output_scale = self.add_weight(
            shape=(1,),
            initializer=tf.keras.initializers.Constant(0.90),
            trainable=True,  # Ensure this is trainable
            constraint=None,  # Remove non-negativity constraint if it's preventing gradient flow
            name='output_scale'
        )
        
        self.output_shift = self.add_weight(
            shape=(1,),
            initializer=tf.keras.initializers.Constant(0.01),
            trainable=True,  # Ensure this is trainable
            constraint=None,
            name='output_shift'
        )
        
        super(MonotonicFunction, self).build(input_shape)
    
    def call(self, inputs):
        water_potential, parameters = inputs
        epsilon = tf.keras.backend.epsilon()
        
        # Handle numerical issues
        water_potential = tf.where(tf.math.is_finite(water_potential), water_potential, tf.zeros_like(water_potential))
        parameters = tf.where(tf.math.is_finite(parameters), parameters, tf.zeros_like(parameters))
        
        # Split parameters
        a_params = parameters[:, :self.num_basis]
        b_params = parameters[:, self.num_basis:]
        
        # Use softplus for positive parameters
        a_params = tf.math.softplus(a_params) + epsilon
        
        # Calculate basis functions
        basis_values = []
        for i in range(self.num_basis):
            # Scale water potential
            scaled_wp = water_potential  # No clipping to allow full range
            
            # Calculate sigmoid logit
            logit = -a_params[:, i:i+1] * scaled_wp + b_params[:, i:i+1]
            g_i = tf.sigmoid(logit)
            
            basis_values.append(g_i)
        
        basis_values = tf.concat(basis_values, axis=-1)
        
        # Calculate weighted sum of basis functions
        weighted_sum = tf.reduce_sum(basis_values * self.basis_weights, axis=-1, keepdims=True)
        
        # Apply sigmoid to get normalized output between 0 and 1
        normalized_output = tf.sigmoid(weighted_sum)
        
        # Scale and shift output to typical water content range (0 to ~0.6)
        water_content = normalized_output * self.output_scale + self.output_shift
        
        return water_content

class ResidualBlock(tf.keras.layers.Layer):
    """
    Residual block with batch normalization and Swish activation
    
    This block implements a residual connection with the following architecture:
    Input -> Dense -> BatchNorm -> Swish -> Dense -> BatchNorm -> Add -> Swish -> Output
                                                                  |
    Input --------> (Optional projection) ----------------------->
    
    Parameters:
    -----------
    units : int
        Number of output units
    projection : bool
        Whether to use projection for input if dimensions don't match
    dropout_rate : float
        Dropout rate applied after the residual connection
    """
    def __init__(self, units, projection=True, dropout_rate=0.1):
        super(ResidualBlock, self).__init__()
        self.units = units
        self.projection = projection
        self.dropout_rate = dropout_rate
        
        # Main path
        self.dense1 = tf.keras.layers.Dense(
            units, 
            kernel_initializer='he_normal'
        )
        self.bn1 = tf.keras.layers.BatchNormalization(epsilon=1e-5)
        self.activation1 = tf.keras.layers.Activation('swish')
        
        self.dense2 = tf.keras.layers.Dense(
            units, 
            kernel_initializer='he_normal'
        )
        self.bn2 = tf.keras.layers.BatchNormalization(epsilon=1e-5)
        
        # Shortcut path (projection if needed)
        self.use_projection = False
        self.projection_layer = None
        self.projection_bn = None
        
        # Final activation and dropout
        self.activation2 = tf.keras.layers.Activation('swish')
        self.dropout = tf.keras.layers.Dropout(dropout_rate)
    
    def build(self, input_shape):
        input_dim = input_shape[-1]
        
        # Create projection layer if input_dim != units
        if input_dim != self.units and self.projection:
            self.use_projection = True
            self.projection_layer = tf.keras.layers.Dense(
                self.units, 
                kernel_initializer='he_normal'
            )
            self.projection_bn = tf.keras.layers.BatchNormalization(epsilon=1e-5)
        
        super(ResidualBlock, self).build(input_shape)
    
    def call(self, inputs, training=False):
        # Main path
        x = self.dense1(inputs)
        x = self.bn1(x, training=training)
        x = self.activation1(x)
        
        x = self.dense2(x)
        x = self.bn2(x, training=training)
        
        # Shortcut path with optional projection
        if self.use_projection:
            shortcut = self.projection_layer(inputs)
            shortcut = self.projection_bn(shortcut, training=training)
        else:
            shortcut = inputs
        
        # Combine paths with residual connection
        x = x + shortcut
        
        # Final activation and dropout
        x = self.activation2(x)
        x = self.dropout(x, training=training)
        
        return x

class SwishActivation(tf.keras.layers.Layer):
    """
    Swish activation function layer: x * sigmoid(x)
    Can be used as a drop-in replacement for other activation layers
    """
    def __init__(self, **kwargs):
        super(SwishActivation, self).__init__(**kwargs)
    
    def call(self, inputs):
        return inputs * tf.sigmoid(inputs)
    
    def compute_output_shape(self, input_shape):
        return input_shape

class EnhancedPropertyEncoder(tf.keras.layers.Layer):
    """
    Enhanced property encoder with residual connections and batch normalization
    
    Architecture:
    Input -> Dense -> BN -> Swish -> ResidualBlock -> ResidualBlock -> Output
    
    Parameters:
    -----------
    embedding_dim : int
        Final embedding dimension
    input_dim : int
        Input dimension (number of features for the property)
    dropout_rate : float
        Dropout rate to apply in residual blocks
    """
    def __init__(self, embedding_dim, input_dim, dropout_rate=0.1, **kwargs):
        super(EnhancedPropertyEncoder, self).__init__(**kwargs)
        self.embedding_dim = embedding_dim
        self.input_dim = input_dim
        self.dropout_rate = dropout_rate
        
        # Initial dense layer to project to intermediate dimension
        self.initial_dense = tf.keras.layers.Dense(
            embedding_dim // 2,
            kernel_initializer='he_normal'
        )
        self.initial_bn = tf.keras.layers.BatchNormalization(epsilon=1e-5)
        self.initial_activation = tf.keras.layers.Activation('swish')
        
        # First residual block (intermediate -> intermediate)
        self.residual1 = ResidualBlock(
            units=embedding_dim // 2,
            projection=False,  # No need for projection here
            dropout_rate=dropout_rate
        )
        
        # Second residual block (intermediate -> final)
        self.residual2 = ResidualBlock(
            units=embedding_dim,
            projection=True,  # Need projection here (dim increase)
            dropout_rate=dropout_rate
        )
    
    def build(self, input_shape):
        # Build the sublayers with the correct input shapes
        self.initial_dense.build(input_shape)
        
        # Get the output shape of the initial dense layer
        intermediate_shape = (input_shape[0], self.embedding_dim // 2)
        
        # Build the batch norm and activation layers
        self.initial_bn.build(intermediate_shape)
        self.initial_activation.build(intermediate_shape)
        
        # Build the first residual block
        self.residual1.build(intermediate_shape)
        
        # Get the output shape of the first residual block
        # (Should be the same as intermediate_shape since projection=False)
        res1_output_shape = intermediate_shape
        
        # Build the second residual block
        self.residual2.build(res1_output_shape)
        
        # Mark the layer as built
        self.built = True
        
        super(EnhancedPropertyEncoder, self).build(input_shape)

    def call(self, inputs, training=False):
        # Initial dense projection
        x = self.initial_dense(inputs)
        x = self.initial_bn(x, training=training)
        x = self.initial_activation(x)
        
        # Apply residual blocks
        x = self.residual1(x, training=training)
        x = self.residual2(x, training=training)
        
        return x
    
    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.embedding_dim)

class CrossAttentionLayer(tf.keras.layers.Layer):
    """
    Specialized cross-attention between soil properties with attention weight return capability
    """
    def __init__(self, embedding_dim, num_heads=4, dropout_rate=0.1):
        super(CrossAttentionLayer, self).__init__()
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        self.dropout_rate = dropout_rate
        
        # Placeholder for attention layer
        self._attention = None
    
    def build(self, input_shape):
        # Ensure input_shape is a list with two elements
        if not isinstance(input_shape, list) or len(input_shape) != 2:
            raise ValueError("Inputs must be a list of two tensors")
        
        # Create MultiHeadAttention layer
        self._attention = tf.keras.layers.MultiHeadAttention(
            num_heads=self.num_heads,
            key_dim=self.embedding_dim // self.num_heads,
            dropout=self.dropout_rate,
            kernel_initializer='glorot_uniform'
        )
        
        # Normalization and processing layers
        self.layer_norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layer_norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        
        # Feed-forward network for processing
        self.ffn = tf.keras.Sequential([
            tf.keras.layers.Dense(self.embedding_dim * 2, activation='swish'),
            tf.keras.layers.Dropout(self.dropout_rate),
            tf.keras.layers.Dense(self.embedding_dim)
        ])
        
        super(CrossAttentionLayer, self).build(input_shape)
    
    def call(self, inputs, training=False, return_attention_weights=False):
        # NEW: Handle stacked property pairs for true cross-attention
        if isinstance(inputs, list) and len(inputs) == 2:
            # Convert pair of properties to stacked format
            prop1, prop2 = inputs
            prop1 = tf.ensure_shape(prop1, [None, self.embedding_dim])
            prop2 = tf.ensure_shape(prop2, [None, self.embedding_dim])
            stacked_features = tf.stack([prop1, prop2], axis=1)  # [batch, 2, embed_dim]
        else:
            # Already stacked
            stacked_features = inputs
        
        # Apply attention on the 2-property stack
        if return_attention_weights:
            attention_output, attention_weights = self._attention(
                query=stacked_features,
                key=stacked_features, 
                value=stacked_features,
                training=training,
                return_attention_scores=True
            )
        else:
            attention_output = self._attention(
                query=stacked_features,
                key=stacked_features,
                value=stacked_features,
                training=training
            )
            attention_weights = None
        
        # Apply residual connections and normalization
        x = self.layer_norm1(stacked_features + attention_output)
        ffn_output = self.ffn(x, training=training)
        output = self.layer_norm2(x + ffn_output)
        
        if return_attention_weights:
            return output, attention_weights
        else:
            return output

class GatingMechanism(tf.keras.layers.Layer):
    """
    Gating mechanism to control information flow between properties
    """
    def __init__(self):
        super(GatingMechanism, self).__init__()
    
    def build(self, input_shape):
        # Ensure input_shape is a list with two elements
        if not isinstance(input_shape, list) or len(input_shape) != 2:
            raise ValueError("Inputs must be a list of two tensors")
        
        # Extract dimensions
        self.input_dim = input_shape[0][-1]
        
        # Create gate parameters
        self.gate_weights = self.add_weight(
            shape=(self.input_dim,),
            initializer=tf.keras.initializers.GlorotNormal(),
            trainable=True,
            name='gate_weights'
        )
        
        super(GatingMechanism, self).build(input_shape)
    
    def call(self, inputs):
        # Ensure inputs is a list
        if not isinstance(inputs, list) or len(inputs) != 2:
            raise ValueError("Inputs must be a list of two tensors")
        
        primary_input, secondary_input = inputs
        
        # Ensure tensor shapes match
        primary_input = tf.ensure_shape(primary_input, secondary_input.shape)
        
        # Create gate
        gate = tf.sigmoid(tf.reduce_sum(primary_input * self.gate_weights, axis=-1, keepdims=True))
        
        # Apply gate to secondary input
        gated_output = primary_input + gate * secondary_input
        
        return gated_output

# Register custom Swish activation
tf.keras.utils.get_custom_objects().update({
    'swish': tf.keras.layers.Activation(lambda x: x * tf.sigmoid(x))
})