import tensorflow as tf from tensorflow import keras import numpy as np from habit.model.layers import ( MonotonicFunction, ResidualBlock, EnhancedPropertyEncoder, SwishActivation, CrossAttentionLayer, GatingMechanism ) # Define custom swish activation function def swish(x): return x * tf.keras.backend.sigmoid(x) # Register the custom activation tf.keras.utils.get_custom_objects().update({'swish': tf.keras.layers.Activation(swish)}) class HABIT(keras.Model): def __init__(self, embedding_dim=128, num_heads=4, num_monotonic_basis=16, dropout_rate=0.1): super(HABIT, self).__init__() # Model properties self.embedding_dim = embedding_dim self.num_heads = num_heads # Store num_heads as an instance variable self.num_monotonic_basis = num_monotonic_basis self.dropout_rate = dropout_rate # Enhanced property encoders with residual connections and batch normalization self.texture_encoder = self._build_property_encoder(3) self.bd_encoder = self._build_property_encoder(1) self.oc_encoder = self._build_property_encoder(1) self.ksat_encoder = self._build_property_encoder(1) # Cross-Attention layers for property interactions self.texture_bd_cross_attention = CrossAttentionLayer( embedding_dim=embedding_dim, num_heads=self.num_heads, # Use the instance variable dropout_rate=dropout_rate ) self.texture_oc_cross_attention = CrossAttentionLayer( embedding_dim=embedding_dim, num_heads=max(2, self.num_heads // 2), # Scale down but ensure at least 2 heads dropout_rate=dropout_rate ) # Gating mechanisms for controlled information flow self.texture_bd_gate = GatingMechanism() self.texture_oc_gate = GatingMechanism() # Fusion layers for enhanced property combinations self.texture_bd_fusion = keras.layers.Dense(embedding_dim, activation='swish') self.texture_oc_fusion = keras.layers.Dense(embedding_dim, activation='swish') # Property attention layer self.property_attention = keras.layers.MultiHeadAttention( num_heads=self.num_heads, # Use the instance variable key_dim=self.embedding_dim // self.num_heads, # Scale key_dim by num_heads value_dim=self.embedding_dim // self.num_heads, # Scale value_dim by num_heads dropout=dropout_rate, use_bias=True ) # Water potential attention layer self.wp_attention = keras.layers.MultiHeadAttention( num_heads=self.num_heads, # Use the instance variable key_dim=self.embedding_dim // self.num_heads, # Scale key_dim by num_heads value_dim=self.embedding_dim // self.num_heads, # Scale value_dim by num_heads dropout=dropout_rate, use_bias=True ) # Layer normalization self.property_layer_norm = keras.layers.LayerNormalization(epsilon=1e-6) self.wp_layer_norm = keras.layers.LayerNormalization(epsilon=1e-6) self.fusion_layer_norm = keras.layers.LayerNormalization(epsilon=1e-6) # Water potential embedding layer self.wp_embedding_layer = keras.layers.Dense(self.embedding_dim, activation='swish') # Additional layers self.soil_dense = keras.layers.Dense(self.embedding_dim, activation='swish') self.flatten = keras.layers.Flatten() # Enhanced feature fusion for property integration self.property_fusion = keras.Sequential([ keras.layers.Dense(self.embedding_dim * 2, activation='swish'), keras.layers.BatchNormalization(epsilon=1e-5), keras.layers.Dropout(self.dropout_rate), keras.layers.Dense(self.embedding_dim, activation='swish') ]) # Parameter network and monotonic function self.parameter_network = keras.Sequential([ keras.layers.Dense(self.embedding_dim, activation='swish'), keras.layers.BatchNormalization(epsilon=1e-5), keras.layers.Dense(self.embedding_dim, activation='swish'), keras.layers.BatchNormalization(epsilon=1e-5), keras.layers.Dense(self.num_monotonic_basis * 2) ]) self.monotonic_function = MonotonicFunction(self.num_monotonic_basis) # Storage for attention weights self._property_attention_weights = None self._wp_attention_weights = None self._cross_attention_weights = None # Will store cross-attention weights def build(self, input_shape): # Ensure all layers are built if isinstance(input_shape, list): texture_shape, bd_shape, oc_shape, ksat_shape, prop_mask_shape, wp_shape = input_shape else: # Default input shapes if not provided texture_shape = (None, 3) bd_shape = (None, 1) oc_shape = (None, 1) ksat_shape = (None, 1) prop_mask_shape = (None, 4) wp_shape = (None, None) # Build property encoders self.texture_encoder.build(texture_shape) self.bd_encoder.build(bd_shape) self.oc_encoder.build(oc_shape) self.ksat_encoder.build(ksat_shape) # Build cross-attention and gating layers self.texture_bd_cross_attention.build([texture_shape, bd_shape]) self.texture_oc_cross_attention.build([texture_shape, oc_shape]) # Ensure get_attention_weights method exists if not hasattr(self, 'get_attention_weights'): def get_attention_weights(self): """Return attention weights for interpretation""" return { 'property_attention': getattr(self, '_property_attention_weights', None), 'wp_attention': getattr(self, '_wp_attention_weights', None), 'cross_attention': getattr(self, '_cross_attention_weights', None) } # Bind the method to the instance import types self.get_attention_weights = types.MethodType(get_attention_weights, self) super().build(input_shape) def _build_property_encoder(self, input_dim): """ Build an enhanced property encoder with residual connections and batch normalization """ return EnhancedPropertyEncoder( embedding_dim=self.embedding_dim, input_dim=input_dim, dropout_rate=self.dropout_rate ) def call(self, inputs, training=False): texture, bd, oc, ksat, properties_mask, water_potential = inputs # Store inputs for access in loss functions self.texture_input = texture self.bd_input = bd self.oc_input = oc self.ksat_input = ksat self.properties_mask_input = properties_mask self.wp_input = water_potential # Create mask for valid water potential values # Note: In log10 scale, 1.0 kPa = 0.0, which is valid! # Only mask out actual padding values (use large negative sentinel) wp_mask = tf.cast(tf.greater(water_potential, -99), tf.float32) # Feature extraction with enhanced property encoders texture_features = self.texture_encoder(texture, training=training) bd_features = self.bd_encoder(bd, training=training) # Ensure texture and bulk density features have correct shape texture_features = tf.ensure_shape(texture_features, [None, self.embedding_dim]) bd_features = tf.ensure_shape(bd_features, [None, self.embedding_dim]) # Process optional properties oc_features = self.oc_encoder(oc, training=training) * tf.expand_dims(properties_mask[:, 2], -1) ksat_features = self.ksat_encoder(ksat, training=training) * tf.expand_dims(properties_mask[:, 3], -1) # Safe handling of optional features oc_features_safe = oc_features if oc_features.shape[-1] == self.embedding_dim else \ tf.zeros_like(texture_features, dtype=texture_features.dtype) # ================================================================= # CROSS-ATTENTION: ALWAYS COMPUTE, USE MASKING FOR CONTROL # ================================================================= # Initialize cross-attention weights storage cross_attention_weights = {} # 1. TEXTURE ↔ BD CROSS-ATTENTION bd_availability = tf.expand_dims(properties_mask[:, 1], -1) # Compute true cross-property attention (2x2 matrix) texture_bd_combined, texture_bd_weights = self.texture_bd_cross_attention( inputs=[texture_features, bd_features], training=training, return_attention_weights=True ) # Extract enhanced features from the 2-property output texture_enhanced_raw = texture_bd_combined[:, 0, :] # First property (texture) bd_enhanced_raw = texture_bd_combined[:, 1, :] # Second property (BD) # Apply gating and masking texture_bd_gated = self.texture_bd_gate([texture_features, texture_enhanced_raw]) * bd_availability bd_texture_gated = self.texture_bd_gate([bd_features, bd_enhanced_raw]) * bd_availability # Fuse features texture_enhanced_by_bd_raw = self.texture_bd_fusion(texture_bd_gated) bd_enhanced_by_texture_raw = self.texture_bd_fusion(bd_texture_gated) # When BD not available, fall back to original features texture_enhanced_by_bd = texture_features + (texture_enhanced_by_bd_raw - texture_features) * bd_availability bd_enhanced_by_texture = bd_features + (bd_enhanced_by_texture_raw - bd_features) * bd_availability # Store attention weights cross_attention_weights['texture_bd'] = texture_bd_weights # 2. TEXTURE ↔ OC CROSS-ATTENTION oc_availability = tf.expand_dims(properties_mask[:, 2], -1) # Compute true cross-property attention (2x2 matrix) texture_oc_combined, texture_oc_weights = self.texture_oc_cross_attention( inputs=[texture_features, oc_features_safe], training=training, return_attention_weights=True ) # Extract enhanced features from the 2-property output texture_enhanced_raw_oc = texture_oc_combined[:, 0, :] # First property (texture) oc_enhanced_raw = texture_oc_combined[:, 1, :] # Second property (OC) # Apply gating and masking texture_oc_gated = self.texture_oc_gate([texture_features, texture_enhanced_raw_oc]) * oc_availability oc_texture_gated = self.texture_oc_gate([oc_features_safe, oc_enhanced_raw]) * oc_availability # Fuse features texture_enhanced_by_oc_raw = self.texture_oc_fusion(texture_oc_gated) oc_enhanced_by_texture_raw = self.texture_oc_fusion(oc_texture_gated) # When OC not available, fall back to original features texture_enhanced_by_oc = texture_features + (texture_enhanced_by_oc_raw - texture_features) * oc_availability oc_enhanced_by_texture = oc_features_safe + (oc_enhanced_by_texture_raw - oc_features_safe) * oc_availability # Store attention weights cross_attention_weights['texture_oc'] = texture_oc_weights # Store all cross-attention weights self._cross_attention_weights = cross_attention_weights # ================================================================= # COMBINE ENHANCED FEATURES # ================================================================= # Combine texture enhancements (instead of simple addition, use learned combination) combined_texture = texture_enhanced_by_bd + texture_enhanced_by_oc combined_texture = self.fusion_layer_norm(combined_texture) # Use enhanced versions of all properties for global attention enhanced_bd = bd_enhanced_by_texture enhanced_oc = oc_enhanced_by_texture # Get batch size batch_size = tf.shape(texture)[0] # Stack ENHANCED property features for global attention mechanism # Shape: [batch_size, 4, embedding_dim] property_features = tf.stack([ combined_texture, # Enhanced texture features enhanced_bd, # Enhanced BD features enhanced_oc, # Enhanced OC features ksat_features # Ksat features (no cross-attention yet) ], axis=1) # Create attention mask for properties # 1 for available properties, 0 for unavailable property_attention_mask = tf.expand_dims( tf.concat([ tf.ones_like(properties_mask[:, 0:2]), # Texture & BD always available tf.reshape(properties_mask[:, 2:4], [batch_size, 2]) # OC & Ksat availability ], axis=1), axis=1 ) # Apply property-level attention property_attention_output, property_weights = self.property_attention( query=property_features, key=property_features, value=property_features, attention_mask=property_attention_mask, return_attention_scores=True, training=training ) # Store attention weights self._property_attention_weights = property_weights # Add & normalize property_features = self.property_layer_norm(property_features + property_attention_output) # Flatten and combine property features using enhanced fusion flattened_features = self.flatten(property_features) soil_embedding = self.property_fusion(flattened_features, training=training) # Process water potential points batch_size = tf.shape(water_potential)[0] num_points = tf.shape(water_potential)[1] # Reshape soil embedding for attention soil_embedding_expanded = tf.expand_dims(soil_embedding, axis=1) # Prepare water potential features wp_expanded = tf.expand_dims(water_potential, axis=-1) # Create embedding for water potential values wp_embedding = self.wp_embedding_layer(wp_expanded) # Create attention mask for water potential wp_attention_mask = tf.expand_dims(wp_mask, axis=1) # Apply water potential attention wp_attention_output, wp_weights = self.wp_attention( query=soil_embedding_expanded, key=wp_embedding, value=wp_embedding, attention_mask=wp_attention_mask, return_attention_scores=True, training=training ) # Store attention weights self._wp_attention_weights = wp_weights # Add & normalize soil_wp_embedding = self.wp_layer_norm(soil_embedding_expanded + wp_attention_output) soil_wp_embedding = tf.squeeze(soil_wp_embedding, axis=1) # Expand soil_wp_embedding for each water potential point soil_expanded = tf.repeat(soil_wp_embedding, num_points, axis=0) # Flatten water potential wp_flat = tf.reshape(water_potential, [batch_size * num_points, 1]) # Combine soil and water potential features combined = tf.concat([soil_expanded, wp_flat], axis=1) # Generate parameters using enhanced parameter network parameters = self.parameter_network(combined, training=training) # Apply monotonic function water_content_flat = self.monotonic_function([wp_flat, parameters]) # Reshape back to original dimensions water_content = tf.reshape(water_content_flat, [batch_size, num_points]) # Apply mask to zero out padded values return water_content * wp_mask def get_attention_weights(self): """Return attention weights for interpretation""" return { 'property_attention': self._property_attention_weights, 'wp_attention': self._wp_attention_weights, 'cross_attention': self._cross_attention_weights }