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)) })