Upload 13 files

app.py

@@ -5,6 +5,7 @@ import numpy as np
 from sklearn.metrics import mean_squared_error, r2_score
 import json
 import pickle
+from rcf_prediction import train_predictor_from_latents
 
 def calculate_fc_accuracy(original_fc, reconstructed_fc):
     """
@@ -89,6 +90,7 @@ def gradio_fc_analysis(data_source, latent_dim, nepochs, bsize, use_hf_dataset):
         demographics = results.get('demographics')
         reconstructed_fc = results.get('reconstructed_fc')
         generated_fc = results.get('generated_fc')
+        outcome_measures = results.get('outcome_measures', None)
         
         # Calculate accuracy metrics
         accuracy_metrics = {}
@@ -108,6 +110,37 @@ def gradio_fc_analysis(data_source, latent_dim, nepochs, bsize, use_hf_dataset):
         if latents is not None and demographics is not None:
             latents_path = save_latents(latents, demographics, file_path=f'latents_dim{latent_dim}.pkl')
             print(f"Saved latents to {latents_path}")
+            
+            # Train a predictor model if we have outcome measures
+            predictor_results = None
+            if outcome_measures is not None and 'wab_aq' in outcome_measures:
+                try:
+                    print("Training WAB-AQ prediction model from latent representations...")
+                    wab_scores = np.array(outcome_measures['wab_aq'])
+                    # Filter out any NaN values
+                    valid_indices = ~np.isnan(wab_scores)
+                    if np.sum(valid_indices) > 5:  # Only train with sufficient data
+                        filtered_latents = latents[valid_indices]
+                        filtered_wab = wab_scores[valid_indices]
+                        
+                        # Extract demographic features for the model
+                        filtered_demographics = {}
+                        for key, values in demographics.items():
+                            if isinstance(values, (list, np.ndarray)) and len(values) >= len(valid_indices):
+                                filtered_demographics[key] = np.array(values)[valid_indices]
+                        
+                        # Train the prediction model with cross-validation
+                        predictor_results = train_predictor_from_latents(
+                            filtered_latents, 
+                            filtered_wab,
+                            filtered_demographics,
+                            cv=5,  # 5-fold cross-validation
+                            n_estimators=100,  # Number of trees in Random Forest
+                            prediction_type="regression"
+                        )
+                        print("WAB-AQ prediction model training complete!")
+                except Exception as e:
+                    print(f"Error training prediction model: {str(e)}")
         
         # Prepare status message with accuracy metrics
         if accuracy_metrics:
@@ -118,8 +151,20 @@ def gradio_fc_analysis(data_source, latent_dim, nepochs, bsize, use_hf_dataset):
                      f"• RMSE: {avg['RMSE']:.6f}\n"
                      f"• R²: {avg['R²']:.6f}\n"
                      f"• Correlation: {avg['Correlation']:.6f}\n"
-                     f"• Cosine Similarity: {avg['Cosine Similarity']:.6f}\n\n"
-                     f"Latent representations saved to results/latents_dim{latent_dim}.pkl")
+                     f"• Cosine Similarity: {avg['Cosine Similarity']:.6f}\n\n")
+            
+            # Add prediction model results if available
+            if predictor_results is not None:
+                cv_results = predictor_results.get('cv_results', {})
+                mean_metrics = cv_results.get('mean_metrics', {})
+                if mean_metrics and 'r2' in mean_metrics:
+                    prediction_r2 = mean_metrics.get('r2', 0)
+                    prediction_rmse = mean_metrics.get('rmse', 0)
+                    status += (f"WAB-AQ Prediction Model Performance:\n"
+                              f"• R²: {prediction_r2:.4f}\n"
+                              f"• RMSE: {prediction_rmse:.4f}\n\n")
+            
+            status += f"Latent representations saved to results/latents_dim{latent_dim}.pkl"
         else:
             status = "Analysis complete! VAE model has been trained and demographic relationships analyzed."
     else:
@@ -216,8 +261,9 @@ def create_interface():
         1. **Data Loading**: The system downloads NIfTI files (P01_rs.nii format) from the SreekarB/OSFData dataset
         2. **Preprocessing**: The fMRI data is processed using the Power 264 atlas and converted to functional connectivity (FC) matrices
         3. **VAE Training**: A conditional VAE model learns the latent representation of brain connectivity
-        4. **Analysis**: The system analyzes relationships between latent brain connectivity patterns and demographic variables
-        5. **Visualization**: Results are displayed showing original FC, reconstructed FC, generated FC, and demographic correlations
+        4. **Predictive Modeling**: The system trains a Random Forest regressor on latent features to predict WAB-AQ scores (aphasia severity)
+        5. **Analysis**: The system analyzes relationships between latent brain connectivity patterns and demographic variables
+        6. **Visualization**: Results are displayed showing original FC, reconstructed FC, generated FC, and demographic correlations
         
         Note: This app works with the SreekarB/OSFData dataset that contains NIfTI files and demographic information.
         """)

main.py

@@ -246,6 +246,12 @@ def run_fc_analysis(data_dir="SreekarB/OSFData",
         
         # If requested, return additional data for accuracy calculations
         if return_data:
+            # Create a structured outcome measures dictionary
+            outcome_measures = {
+                'wab_aq': demo_data[3],  # WAB-AQ scores
+                # Could add other outcome measures here
+            }
+            
             results = {
                 'vae': vae,
                 'X': X,
@@ -253,7 +259,8 @@ def run_fc_analysis(data_dir="SreekarB/OSFData",
                 'demographics': demographics,
                 'reconstructed_fc': reconstructed_fc,
                 'generated_fc': generated_fc,
-                'analysis_results': analysis_results
+                'analysis_results': analysis_results,
+                'outcome_measures': outcome_measures
             }
             return fig, results
         

rcf_prediction.py

@@ -54,10 +54,54 @@ class AphasiaTreatmentPredictor:
             tuple: Combined features array and feature names
         """
         if isinstance(demographics, dict):
-            demo_df = pd.DataFrame(demographics)
+            # For dictionary input, ensure all arrays are same length as latents
+            n_samples = latents.shape[0]
+            aligned_demos = {}
+            
+            for key, values in demographics.items():
+                if len(values) != n_samples:
+                    print(f"WARNING: Demographics '{key}' length ({len(values)}) doesn't match latents ({n_samples})")
+                    # Truncate or pad to match latent samples
+                    if len(values) > n_samples:
+                        aligned_demos[key] = values[:n_samples]  # Truncate
+                        print(f"  Truncated '{key}' to {n_samples} samples")
+                    else:
+                        # Pad with repeated values or zeros depending on type
+                        if len(values) > 0:
+                            # Use mean for numerical, mode for categorical
+                            if isinstance(values[0], (int, float, np.number)):
+                                filler = np.mean(values)
+                            else:
+                                # Use most common value
+                                from collections import Counter
+                                filler = Counter(values).most_common(1)[0][0]
+                            
+                            padding = [filler] * (n_samples - len(values))
+                            aligned_demos[key] = list(values) + padding
+                            print(f"  Padded '{key}' with {filler} to {n_samples} samples")
+                        else:
+                            # Empty array, fill with zeros
+                            aligned_demos[key] = [0] * n_samples
+                            print(f"  Filled empty '{key}' with zeros to {n_samples} samples")
+                else:
+                    aligned_demos[key] = values
+            
+            demo_df = pd.DataFrame(aligned_demos)
         else:
             demo_df = demographics.copy()
             
+            # Ensure DataFrame has same number of rows as latents
+            if len(demo_df) != latents.shape[0]:
+                print(f"WARNING: Demographics DataFrame size ({len(demo_df)}) doesn't match latents ({latents.shape[0]})")
+                if len(demo_df) > latents.shape[0]:
+                    demo_df = demo_df.iloc[:latents.shape[0]]  # Truncate
+                    print(f"  Truncated demographics to {latents.shape[0]} samples")
+                else:
+                    # Cannot easily pad DataFrame, use last row or means
+                    print(f"  ERROR: Cannot pad demographics DataFrame - using latents only")
+                    # Create a DataFrame with the same columns but zeros
+                    demo_df = pd.DataFrame(0, index=range(latents.shape[0]), columns=demo_df.columns)
+            
         # Get categorical columns
         cat_columns = demo_df.select_dtypes(include=['object']).columns.tolist()
         
@@ -71,7 +115,16 @@ class AphasiaTreatmentPredictor:
         feature_names = latent_names + demo_names
         
         # Combine latents with demographics
-        features = np.hstack([latents, demo_df.values])
+        try:
+            features = np.hstack([latents, demo_df.values])
+        except ValueError as e:
+            print(f"ERROR combining features: {e}")
+            print(f"Latents shape: {latents.shape}, Demographics shape: {demo_df.values.shape}")
+            # Fall back to using just latents
+            print("Falling back to using only latent features")
+            features = latents
+            feature_names = latent_names
+            
         return features, feature_names
         
     def fit(self, latents, demographics, treatment_outcomes):
@@ -90,6 +143,11 @@ class AphasiaTreatmentPredictor:
         self.feature_names = feature_names
         
         logger.info(f"Training {self.prediction_type} model with {X.shape[0]} samples and {X.shape[1]} features")
+        print(f"Random Forest: Building {self.n_estimators} trees...")
+        
+        # Track progress during fit with verbose
+        # Set verbose to 2 for detailed per-tree progress
+        self.model.verbose = 1
         self.model.fit(X, treatment_outcomes)
         
         # Calculate feature importance
@@ -98,6 +156,7 @@ class AphasiaTreatmentPredictor:
             'importance': self.model.feature_importances_
         }).sort_values('importance', ascending=False)
         
+        print(f"Random Forest: Training complete. Top features: {', '.join(self.feature_importance['feature'].head(3).tolist())}")
         return self
         
     def predict(self, latents, demographics):
@@ -160,31 +219,54 @@ class AphasiaTreatmentPredictor:
         X, feature_names = self.prepare_features(latents, demographics)
         self.feature_names = feature_names
         
-        logger.info(f"Running {n_splits}-fold cross-validation")
-        
-        kf = KFold(n_splits=n_splits, shuffle=True, random_state=self.random_state)
+        # Adjust n_splits if we have too few samples
+        sample_count = len(treatment_outcomes)
+        if sample_count < n_splits * 2:  # Need at least 2 samples per fold
+            adjusted_n_splits = max(2, sample_count // 2)  # At least 2 folds, each with multiple samples
+            logger.warning(f"Too few samples ({sample_count}) for {n_splits} folds. Adjusting to {adjusted_n_splits} folds.")
+            print(f"Random Forest: Starting {adjusted_n_splits}-fold cross-validation with {sample_count} samples")
+            n_splits = adjusted_n_splits
+        else:
+            logger.info(f"Running {n_splits}-fold cross-validation on {sample_count} samples")
+            print(f"Random Forest: Starting {n_splits}-fold cross-validation with {sample_count} samples")
+        
+        # Use stratified KFold for regression to ensure balanced folds
+        # or LeaveOneOut for very small datasets
+        if sample_count <= 5:
+            from sklearn.model_selection import LeaveOneOut
+            logger.warning(f"Using Leave-One-Out CV for small dataset with {sample_count} samples")
+            print(f"Random Forest: Using Leave-One-Out cross-validation due to small sample size ({sample_count})")
+            kf = LeaveOneOut()
+            cv_iterator = kf.split(X)
+        else:
+            kf = KFold(n_splits=n_splits, shuffle=True, random_state=self.random_state)
+            cv_iterator = kf.split(X)
         
         cv_scores = []
         predictions = np.zeros_like(treatment_outcomes)
         prediction_stds = np.zeros_like(treatment_outcomes)
         fold_metrics = []
         
-        for fold, (train_idx, test_idx) in enumerate(kf.split(X)):
+        for fold, (train_idx, test_idx) in enumerate(cv_iterator):
             X_train, X_test = X[train_idx], X[test_idx]
             y_train, y_test = treatment_outcomes[train_idx], treatment_outcomes[test_idx]
             
+            print(f"Random Forest: Training fold {fold+1}/{n_splits} - {len(X_train)} training samples, {len(X_test)} test samples")
+            
             # Clone the model for this fold
             if self.prediction_type == "classification":
                 fold_model = RandomForestClassifier(
                     n_estimators=self.n_estimators,
                     max_depth=self.max_depth,
-                    random_state=self.random_state
+                    random_state=self.random_state,
+                    verbose=1  # Add verbosity
                 )
             else:
                 fold_model = RandomForestRegressor(
                     n_estimators=self.n_estimators,
                     max_depth=self.max_depth,
-                    random_state=self.random_state
+                    random_state=self.random_state,
+                    verbose=1  # Add verbosity
                 )
             
             # Train the model
@@ -199,13 +281,34 @@ class AphasiaTreatmentPredictor:
             # Calculate metrics
             if self.prediction_type == "regression":
                 rmse = np.sqrt(mean_squared_error(y_test, pred))
-                r2 = r2_score(y_test, pred)
+                
+                # R-squared requires at least 2 samples and some variance in the target
+                if len(y_test) >= 2 and np.var(y_test) > 1e-10:
+                    r2 = r2_score(y_test, pred)
+                else:
+                    r2 = np.nan
+                    logger.warning(f"Fold {fold+1}: R² not calculated (insufficient samples or variance)")
+                    print(f"Random Forest: Fold {fold+1} - R² not calculated (insufficient samples or variance)")
+                
+                # MSE can always be calculated
+                mse = rmse**2
+                
                 metrics = {
                     "r2": r2,
                     "rmse": rmse,
-                    "mse": rmse**2
+                    "mse": mse
                 }
                 
+                # Add other useful metrics if there are enough samples
+                if len(y_test) >= 2 and np.var(y_test) > 1e-10:
+                    from sklearn.metrics import explained_variance_score
+                    try:
+                        ev = explained_variance_score(y_test, pred)
+                        metrics["explained_variance"] = ev
+                    except:
+                        # Skip if it can't be calculated
+                        pass
+                
                 # Get prediction intervals using tree variance
                 tree_predictions = np.array([tree.predict(X_test) 
                                           for tree in fold_model.estimators_])
@@ -233,14 +336,43 @@ class AphasiaTreatmentPredictor:
             fold_metrics.append(metrics)
             logger.info(f"Fold {fold+1} metrics: {metrics}")
             
+            # Print a more user-friendly version of the fold results
+            if self.prediction_type == "regression":
+                r2_val = metrics.get('r2', np.nan)
+                rmse_val = metrics.get('rmse', np.nan)
+                r2_text = f"R² = {r2_val:.4f}" if not np.isnan(r2_val) else "R² = N/A"
+                print(f"Random Forest: Fold {fold+1} results - {r2_text}, RMSE = {rmse_val:.4f}")
+            else:
+                acc_val = metrics.get('accuracy', 0)
+                f1_val = metrics.get('f1', 0)
+                print(f"Random Forest: Fold {fold+1} results - Accuracy = {acc_val:.4f}, F1 = {f1_val:.4f}")
+            
         # Calculate average metrics
         avg_metrics = {}
         for key in fold_metrics[0].keys():
-            avg_metrics[key] = np.mean([fold[key] for fold in fold_metrics])
+            # Filter out nan values when calculating means
+            values = [fold[key] for fold in fold_metrics if key in fold and not (isinstance(fold[key], float) and np.isnan(fold[key]))]
+            if values:  # Only calculate mean if we have valid values
+                avg_metrics[key] = np.mean(values)
+            else:
+                avg_metrics[key] = np.nan
             
         logger.info(f"Average CV metrics: {avg_metrics}")
         
+        # Print a summary of cross-validation performance
+        if self.prediction_type == "regression":
+            r2_avg = avg_metrics.get('r2', np.nan)
+            rmse_avg = avg_metrics.get('rmse', np.nan)
+            r2_text = f"R² = {r2_avg:.4f}" if not np.isnan(r2_avg) else "R² = N/A"
+            print(f"Random Forest: Cross-validation complete - Average {r2_text}, RMSE = {rmse_avg:.4f}")
+        else:
+            acc_avg = avg_metrics.get('accuracy', 0)
+            f1_avg = avg_metrics.get('f1', 0)
+            print(f"Random Forest: Cross-validation complete - Average Accuracy = {acc_avg:.4f}, F1 = {f1_avg:.4f}")
+        
         # Train final model on all data
+        print(f"Random Forest: Training final model on all {len(X)} samples...")
+        self.model.verbose = 1
         self.model.fit(X, treatment_outcomes)
         
         # Calculate feature importance