Merge pull request #7 from RaykKretzschmar/copilot/sub-pr-3-yet-again

Align DummySpectralCNN test model with real SpectralCNN tuple return signature

Author: RaykKretzschmar

spectra2prop/utils/explainability.py

@@ -79,14 +79,20 @@ def generate_heatmap(self, input_tensor, target_class_idx=None, normalize=True):
         # 1. Forward pass
         # The hooks will capture self.activations here
         output = self.model(input_tensor)
+        
+        # Handle tuple output (logits, embedding) from SpectralCNN
+        if isinstance(output, tuple):
+            logits = output[0]
+        else:
+            logits = output
 
         # Determine target class
         if target_class_idx is None:
-            target_class_idx = output.argmax(dim=1).item()
+            target_class_idx = logits.argmax(dim=1).item()
 
         # 2. Backward pass
         # Create a one-hot target tensor to backpropagate specific class score
-        target_score = output[:, target_class_idx]
+        target_score = logits[:, target_class_idx]
         target_score.backward(retain_graph=True)
 
         # 3. Compute Weights

tests/test_explainability.py

@@ -25,8 +25,9 @@ def forward(self, x):
         x = self.relu(self.conv2(x))
         x = self.gap(x)
         x = x.view(x.size(0), -1)
-        x = self.fc(x)
-        return x
+        embedding = x  # Save features before final layer
+        logits = self.fc(x)
+        return logits, embedding  # Return tuple (logits, features) to match real SpectralCNN model interface
 
 class NoConvModel(nn.Module):
     """A model with no Conv1d layers to test error handling."""
@@ -107,6 +108,8 @@ def test_heatmap_automatic_class_selection(self):
         # Get model output to determine expected class
         with torch.no_grad():
             output = self.model(self.input_tensor)
+            if isinstance(output, tuple):
+                output = output[0]
             expected_class = output.argmax(dim=1).item()
 
         # Generate heatmap with automatic class selection