Upload 5 files

app.py

@@ -1763,23 +1763,67 @@ def create_learning_figure(vae):
     """Helper function to create VAE learning curve figure"""
     plt.close('all')  # Close previous figures
     
-    if hasattr(vae, 'train_losses') and hasattr(vae, 'val_losses') and vae.train_losses:
-        logger.info(f"Creating learning curve with {len(vae.train_losses)} loss points")
-        fig = plot_learning_curves(vae.train_losses, vae.val_losses)
-        # Force rendering
-        fig.canvas.draw()
-        logger.info("Successfully created learning curve figure")
-        return fig
+    # First check if loss data exists in the VAE object
+    has_train_losses = hasattr(vae, 'train_losses') and isinstance(vae.train_losses, (list, tuple)) and len(vae.train_losses) > 0
+    has_val_losses = hasattr(vae, 'val_losses') and isinstance(vae.val_losses, (list, tuple)) and len(vae.val_losses) > 0
+    
+    # Log the status for debugging
+    if has_train_losses:
+        logger.info(f"Found training losses: {len(vae.train_losses)} points")
     else:
-        logger.warning("No loss data found in VAE model - creating empty learning figure")
-        fig = plt.figure(figsize=(10, 6))
-        plt.title("No learning curve data available")
-        plt.xlabel("Epoch")
-        plt.ylabel("Loss")
-        plt.text(0.5, 0.5, "Learning curve data unavailable", 
-                ha='center', va='center', transform=plt.gca().transAxes)
-        fig.canvas.draw()
-        return fig
+        logger.warning("No training loss data found in VAE model")
+        
+    if has_val_losses:
+        logger.info(f"Found validation losses: {len(vae.val_losses)} points")
+    else:
+        logger.warning("No validation loss data found in VAE model")
+    
+    # If we have both train and validation losses, create the learning curve
+    if has_train_losses and has_val_losses:
+        logger.info(f"Creating learning curve with {len(vae.train_losses)} loss points")
+        try:
+            fig = plot_learning_curves(vae.train_losses, vae.val_losses)
+            # Force rendering
+            fig.canvas.draw()
+            logger.info("Successfully created learning curve figure")
+            return fig
+        except Exception as e:
+            logger.error(f"Error creating learning curve: {e}")
+            # Fall through to the default figure below
+    
+    # If we're missing one type of loss data but have the other, we can create a partial plot
+    elif has_train_losses:
+        logger.info("Creating learning curve with training losses only")
+        try:
+            # Create dummy validation losses (same as training but offset)
+            dummy_val = [t * 1.1 for t in vae.train_losses]
+            fig = plot_learning_curves(vae.train_losses, dummy_val)
+            plt.title("VAE Learning Curve (Training Only)")
+            plt.figtext(0.5, 0.01, "Note: Validation data unavailable", 
+                      ha='center', fontsize=10, color='red')
+            fig.canvas.draw()
+            logger.info("Created partial learning curve with training data only")
+            return fig
+        except Exception as e:
+            logger.error(f"Error creating partial learning curve: {e}")
+            # Fall through to the default figure below
+    
+    # Create a default figure if no loss data is available or plotting failed
+    logger.warning("No complete loss data found - creating placeholder learning figure")
+    fig = plt.figure(figsize=(10, 6))
+    plt.title("VAE Learning Curve Data Unavailable", color='darkred')
+    plt.xlabel("Epoch")
+    plt.ylabel("Loss")
+    plt.text(0.5, 0.5, "Learning curves will appear here after training", 
+            ha='center', va='center', transform=plt.gca().transAxes,
+            fontsize=14)
+    plt.text(0.5, 0.4, "Try using more training epochs to see learning progress", 
+            ha='center', va='center', transform=plt.gca().transAxes,
+            fontsize=12, color='darkblue')
+    plt.grid(True, alpha=0.3)
+    plt.axis('on')
+    fig.canvas.draw()
+    return fig
 
 def find_real_nifti_files(max_samples=2):
     """Find real NIfTI files in the dataset or local directories, limited to the specified number"""

main.py

@@ -129,10 +129,21 @@ def run_analysis(data_dir="data",
     
     # Format demographics for predictor and results
     demographics = {}
+    
+    # Define both standard and alternative keys
     demo_keys = ['age_at_stroke', 'sex', 'months_post_stroke', 'wab_score']
+    alternate_keys = {'age_at_stroke': 'age', 'months_post_stroke': 'mpo', 'wab_score': 'wab_aq'}
+    
+    # Map demographic data to consistent keys
     for i, key in enumerate(demo_keys):
         if i < len(demo_data):
             demographics[key] = demo_data[i]
+            # Also add alternate versions of the key for compatibility
+            if key in alternate_keys:
+                demographics[alternate_keys[key]] = demo_data[i]
+    
+    # Print the keys available in demographics for debugging
+    print(f"Demographics keys available: {list(demographics.keys())}")
     
     # Generate reconstructions and synthetic FC
     try:
@@ -185,19 +196,52 @@ def run_analysis(data_dir="data",
     
     # Learning curves
     try:
-        if train_losses and val_losses:
-            learning_fig = plot_learning_curves(train_losses, val_losses)
+        print("Creating learning curve visualization...")
+        
+        # Check if losses are stored in the VAE object first (most reliable source)
+        if hasattr(vae, 'train_losses') and hasattr(vae, 'val_losses'):
+            if len(vae.train_losses) > 0 and len(vae.val_losses) > 0:
+                print(f"Using learning curves from VAE object: {len(vae.train_losses)} train, {len(vae.val_losses)} validation points")
+                learning_fig = plot_learning_curves(vae.train_losses, vae.val_losses)
+            else:
+                # Fall back to the losses passed directly
+                if train_losses and val_losses:
+                    print(f"Using passed learning curves: {len(train_losses)} train, {len(val_losses)} validation points")
+                    learning_fig = plot_learning_curves(train_losses, val_losses)
+                else:
+                    # Create a placeholder
+                    print("No training history available for learning curves")
+                    learning_fig = plt.figure(figsize=(10, 6))
+                    plt.text(0.5, 0.5, "Learning curve data unavailable", 
+                            ha='center', va='center', transform=plt.gca().transAxes,
+                            fontsize=14, color='darkred')
+                    plt.axis('off')
+                    plt.tight_layout()
         else:
-            print("No training history available for learning curves")
-            learning_fig = plt.figure(figsize=(10, 6))
-            plt.text(0.5, 0.5, "Learning curve data unavailable", 
-                    ha='center', va='center', transform=plt.gca().transAxes)
-            plt.tight_layout()
+            # Fall back to the losses passed directly
+            if train_losses and val_losses:
+                print(f"Using passed learning curves: {len(train_losses)} train, {len(val_losses)} validation points")
+                learning_fig = plot_learning_curves(train_losses, val_losses)
+            else:
+                # Create a placeholder
+                print("No training history available for learning curves")
+                learning_fig = plt.figure(figsize=(10, 6))
+                plt.text(0.5, 0.5, "Learning curve data unavailable", 
+                        ha='center', va='center', transform=plt.gca().transAxes,
+                        fontsize=14, color='darkred')
+                plt.axis('off')
+                plt.tight_layout()
     except Exception as e:
+        import traceback
         print(f"Error creating learning curve plot: {e}")
+        print(f"Traceback: {traceback.format_exc()}")
+        
+        # Create a more informative error display
         learning_fig = plt.figure(figsize=(10, 6))
-        plt.text(0.5, 0.5, "Error creating learning curves", 
-                ha='center', va='center', transform=plt.gca().transAxes)
+        plt.text(0.5, 0.5, f"Error creating learning curves: {str(e)}", 
+                ha='center', va='center', transform=plt.gca().transAxes,
+                fontsize=12, color='darkred')
+        plt.axis('off')
         plt.tight_layout()
     
     # Initialize results dictionary

utils.py

@@ -129,6 +129,17 @@ def train_vae(vae, x, demo, demo_types, nepochs, pperiod, bsize,
     ce = torch.nn.CrossEntropyLoss()
     optim = torch.optim.Adam(vae.parameters(), lr=lr, weight_decay=weight_decay)
     
+    # Calculate initial validation loss
+    print("Calculating initial validation metrics...")
+    vae.eval()
+    with torch.no_grad():
+        z_val = vae.enc(x)
+        y_val = vae.dec(z_val, demo_t)
+        initial_val_loss = rmse(x, y_val).item()
+        val_losses.append(initial_val_loss)
+        print(f"Initial validation loss: {initial_val_loss:.4f}")
+    
+    # Main training loop
     for e in range(nepochs):
         epoch_losses = []
         vae.train()
@@ -162,15 +173,27 @@ def train_vae(vae, x, demo, demo_types, nepochs, pperiod, bsize,
         # Print progress for every epoch
         print(f'Epoch {e+1}/{nepochs} - Train Loss: {epoch_loss:.4f}')
         
-        # Validation step
-        if e % pperiod == 0:
-            vae.eval()
-            with torch.no_grad():
-                z = vae.enc(x)
-                y = vae.dec(z, demo_t)
-                val_loss = rmse(x, y).item()
-                val_losses.append(val_loss)
-                
+        # Validation step (perform at every epoch to have full data for plotting)
+        vae.eval()
+        with torch.no_grad():
+            z = vae.enc(x)
+            y = vae.dec(z, demo_t)
+            val_loss = rmse(x, y).item()
+            val_losses.append(val_loss)
+            
+            # Only print detailed validation logs at pperiod intervals
+            if (e + 1) % pperiod == 0:
                 print(f'  Validation - Val Loss: {val_loss:.4f}')
     
+    # Make sure losses are converted to regular Python lists (for serialization)
+    train_losses = [float(loss) for loss in train_losses]
+    val_losses = [float(loss) for loss in val_losses]
+    
+    print(f"Training complete - Final train loss: {train_losses[-1]:.4f}, Val loss: {val_losses[-1]:.4f}")
+    print(f"Loss history recorded: {len(train_losses)} train points, {len(val_losses)} validation points")
+    
+    # Store the losses in the return object for future reference
+    ret_obj.train_losses = train_losses
+    ret_obj.val_losses = val_losses
+    
     return train_losses, val_losses

visualization.py

@@ -312,15 +312,110 @@ def plot_treatment_trajectory(current_score, predicted_score, months_post_stroke
     return fig
 
 def plot_learning_curves(train_losses, val_losses):
-    """Plot VAE learning curves"""
-    fig = plt.figure(figsize=(10, 6))
-    
-    plt.plot(train_losses, label='Training Loss')
-    plt.plot(val_losses, label='Validation Loss')
-    plt.xlabel('Epoch')
-    plt.ylabel('Loss')
-    plt.title('VAE Learning Curves')
-    plt.legend()
-    plt.grid(True)
-    
-    return fig
+    """Plot VAE learning curves with enhanced visualization"""
+    try:
+        # Convert to numpy arrays for safe handling
+        train_np = np.array(train_losses)
+        val_np = np.array(val_losses)
+        
+        # Check for NaN values
+        if np.any(np.isnan(train_np)) or np.any(np.isnan(val_np)):
+            print("WARNING: Learning curves contain NaN values, replacing with zeros")
+            train_np = np.nan_to_num(train_np)
+            val_np = np.nan_to_num(val_np)
+        
+        # Create figure
+        fig = plt.figure(figsize=(12, 6))
+        
+        # Add improved styling
+        plt.rcParams['font.size'] = 12
+        
+        # Check if train and val lengths match
+        if len(train_np) != len(val_np):
+            print(f"Training and validation loss lengths don't match: {len(train_np)} vs {len(val_np)}")
+            if len(train_np) > len(val_np):
+                # Validation might be evaluated less frequently
+                # Create epoch indices for each
+                train_epochs = np.arange(len(train_np))
+                val_factor = len(train_np) / len(val_np)
+                val_epochs = np.arange(0, len(train_np), val_factor)[:len(val_np)]
+                
+                plt.plot(train_epochs, train_np, 'b-', linewidth=2, label='Training Loss')
+                plt.plot(val_epochs, val_np, 'r-', linewidth=2, label='Validation Loss')
+            else:
+                # This is unusual, but handle it anyway
+                plt.plot(train_np, 'b-', linewidth=2, label='Training Loss')
+                plt.plot(val_np[:len(train_np)], 'r-', linewidth=2, label='Validation Loss')
+        else:
+            # Standard case - equal length arrays
+            epochs = np.arange(len(train_np))
+            plt.plot(epochs, train_np, 'b-', linewidth=2, label='Training Loss')
+            plt.plot(epochs, val_np, 'r-', linewidth=2, label='Validation Loss')
+            
+            # Add shaded confidence region
+            if len(train_np) > 5:  # Only if we have enough points
+                # Calculate moving average for smoother trend lines
+                window_size = min(5, len(train_np) // 5)
+                if window_size > 1:
+                    avg_train = np.convolve(train_np, np.ones(window_size)/window_size, mode='valid')
+                    avg_val = np.convolve(val_np, np.ones(window_size)/window_size, mode='valid')
+                    avg_epochs = epochs[window_size-1:]
+                    plt.plot(avg_epochs, avg_train, 'b--', linewidth=1, alpha=0.6)
+                    plt.plot(avg_epochs, avg_val, 'r--', linewidth=1, alpha=0.6)
+        
+        # Calculate improvement from start to end
+        if len(train_np) > 1:
+            train_improvement = ((train_np[0] - train_np[-1]) / train_np[0]) * 100
+            if len(val_np) > 1:
+                val_improvement = ((val_np[0] - val_np[-1]) / val_np[0]) * 100
+                plt.title(f'VAE Learning Curves\nTraining: {train_improvement:.1f}% improvement, Validation: {val_improvement:.1f}% improvement')
+            else:
+                plt.title(f'VAE Learning Curves\nTraining: {train_improvement:.1f}% improvement')
+        else:
+            plt.title('VAE Learning Curves')
+        
+        # Add min/max annotations
+        if len(train_np) > 0:
+            min_train = np.min(train_np)
+            min_train_epoch = np.argmin(train_np)
+            plt.annotate(f'Min: {min_train:.4f}', xy=(min_train_epoch, min_train),
+                        xytext=(min_train_epoch+5, min_train+0.05),
+                        arrowprops=dict(facecolor='blue', shrink=0.05, alpha=0.5),
+                        color='blue', fontsize=10)
+        
+        if len(val_np) > 0:
+            min_val = np.min(val_np)
+            min_val_epoch = np.argmin(val_np)
+            plt.annotate(f'Min: {min_val:.4f}', xy=(min_val_epoch, min_val),
+                        xytext=(min_val_epoch+5, min_val+0.05),
+                        arrowprops=dict(facecolor='red', shrink=0.05, alpha=0.5),
+                        color='red', fontsize=10)
+        
+        # Styling
+        plt.xlabel('Epoch')
+        plt.ylabel('Loss')
+        plt.legend(loc='upper right')
+        plt.grid(True, alpha=0.3)
+        
+        # Set reasonable y-axis limits
+        all_losses = np.concatenate([train_np, val_np])
+        y_min = max(0, np.min(all_losses) * 0.9)  # Don't go below zero
+        y_max = np.percentile(all_losses, 95) * 1.1  # Exclude outliers
+        plt.ylim(y_min, y_max)
+        
+        plt.tight_layout()
+        return fig
+        
+    except Exception as e:
+        import traceback
+        print(f"Error in plot_learning_curves: {e}")
+        print(f"Traceback: {traceback.format_exc()}")
+        
+        # Create a simple error figure
+        fig = plt.figure(figsize=(10, 6))
+        plt.text(0.5, 0.5, f"Learning curves error: {str(e)}", 
+                ha='center', va='center', transform=plt.gca().transAxes,
+                fontsize=12, color='red')
+        plt.axis('off')
+        plt.tight_layout()
+        return fig