Upload 10 files

app.py

@@ -1,3 +1,10 @@
+import os
+import sys
+
+# Set Huggingface cache directory to avoid permission issues
+os.environ['TRANSFORMERS_CACHE'] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'hf_cache')
+os.makedirs(os.environ['TRANSFORMERS_CACHE'], exist_ok=True)
+
 import gradio as gr
 from main import run_analysis
 from rcf_prediction import AphasiaTreatmentPredictor
@@ -10,7 +17,6 @@ matplotlib.rcParams['savefig.dpi'] = 100
 import matplotlib.pyplot as plt
 from data_preprocessing import preprocess_fmri_to_fc, process_single_fmri
 from visualization import plot_fc_matrices, plot_learning_curves
-import os
 import glob
 from sklearn.metrics import mean_squared_error, r2_score
 import json

config.py

@@ -1,8 +1,8 @@
 # Model configuration
 MODEL_CONFIG = {
     'latent_dim': 32,
-    'nepochs': 1000,
-    'bsize': 16,
+    'nepochs': 100,     # Changed from 1000 to 100 for faster testing
+    'bsize': 5,         # Changed from 16 to 5 for small sample sizes
     'loss_rec_mult': 100,
     'loss_decor_mult': 10,
     'lr': 1e-4

hf_cache/version.txt

@@ -0,0 +1 @@
+1

main.py

@@ -3,6 +3,11 @@ import numpy as np  # Make sure numpy is imported at the top level
 import torch
 from pathlib import Path
 import pandas as pd
+
+# Set Huggingface cache directory to avoid permission issues
+os.environ['TRANSFORMERS_CACHE'] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'hf_cache')
+os.makedirs(os.environ['TRANSFORMERS_CACHE'], exist_ok=True)
+
 from data_preprocessing import load_and_preprocess_data
 from vae_model import DemoVAE
 from rcf_prediction import AphasiaTreatmentPredictor

test_small_sample.py

@@ -0,0 +1,77 @@
+import os
+# Set Huggingface cache directory to avoid permission issues
+os.environ['TRANSFORMERS_CACHE'] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'hf_cache')
+os.makedirs(os.environ['TRANSFORMERS_CACHE'], exist_ok=True)
+
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from vae_model import DemoVAE
+from visualization import plot_learning_curves, plot_fc_matrices
+from config import MODEL_CONFIG
+
+# Create small synthetic dataset with only 5 samples
+input_dim = 100
+n_samples = 5
+demo_dim = 4
+
+print(f"Creating test dataset with {n_samples} samples...")
+
+# Synthetic FC matrices (Upper triangular values)
+X = np.random.randn(n_samples, input_dim)
+
+# Synthetic demographics
+demo_data = [
+    np.random.normal(60, 10, n_samples),  # age
+    np.random.choice(['M', 'F'], n_samples),  # sex
+    np.random.normal(24, 12, n_samples),  # months post stroke
+    np.random.normal(50, 15, n_samples)   # WAB score
+]
+
+# Types of demographics
+demo_types = ['continuous', 'categorical', 'continuous', 'continuous']
+
+# Initialize model with updated config
+print("Config settings:")
+print(f"- Epochs: {MODEL_CONFIG['nepochs']}")
+print(f"- Batch size: {MODEL_CONFIG['bsize']}")
+print(f"- Latent dim: {MODEL_CONFIG['latent_dim']}")
+
+print("Initializing model...")
+vae = DemoVAE(**MODEL_CONFIG)
+
+# Train model
+print(f"Training model with {n_samples} samples...")
+train_losses, val_losses = vae.fit(X, demo_data, demo_types)
+
+print(f"Training complete! Final train loss: {train_losses[-1]:.4f}")
+print(f"Final validation loss: {val_losses[-1]:.4f}")
+
+# Save model
+os.makedirs("models", exist_ok=True)
+os.makedirs("results", exist_ok=True)
+print("Saving model...")
+vae.save('models/vae_model_small.pt')
+
+# Create learning curve visualization
+print("Generating learning curve visualization...")
+learning_fig = plot_learning_curves(train_losses, val_losses)
+learning_fig.savefig('results/learning_curves_small.png')
+print("Learning curve saved to results/learning_curves_small.png")
+
+# Generate reconstructed data
+print("Generating reconstructions...")
+reconstructed = vae.transform(X, demo_data, demo_types)
+
+# Get a single sample for FC visualization
+original = X[0].reshape(10, 10)  # Reshape to square matrix for visualization
+recon = reconstructed[0].reshape(10, 10)
+generated = vae.transform(1, [d[:1] for d in demo_data], demo_types)[0].reshape(10, 10)
+
+# Create FC visualization
+print("Generating FC matrix visualization...")
+fc_fig = plot_fc_matrices(original, recon, generated)
+fc_fig.savefig('results/fc_visualization_small.png')
+print("FC visualization saved to results/fc_visualization_small.png")
+
+print("Test with small sample size completed successfully!")

test_train.py

@@ -0,0 +1,59 @@
+import numpy as np
+import torch
+from vae_model import DemoVAE
+
+# Create synthetic data
+input_dim = 100
+n_samples = 20
+demo_dim = 4
+
+# Synthetic FC matrices (Upper triangular values)
+X = np.random.randn(n_samples, input_dim)
+
+# Synthetic demographics
+demo_data = [
+    np.random.normal(60, 10, n_samples),  # age
+    np.random.choice([0, 1], n_samples),  # sex
+    np.random.normal(24, 12, n_samples),  # months post stroke
+    np.random.normal(50, 15, n_samples)   # WAB score
+]
+
+# Types of demographics
+demo_types = ['continuous', 'categorical', 'continuous', 'continuous']
+
+# Initialize model
+model_config = {
+    'latent_dim': 16,
+    'nepochs': 5,
+    'bsize': 5,
+    'use_cuda': False
+}
+
+print("Initializing model...")
+vae = DemoVAE(**model_config)
+
+# Train model
+print("Training model...")
+train_losses, val_losses = vae.fit(X, demo_data, demo_types)
+
+print(f"Training complete! Train loss: {train_losses[-1]}, Val loss: {val_losses[-1]}")
+
+# Check shapes of losses
+print(f"Train losses shape: {len(train_losses)}")
+print(f"Val losses shape: {len(val_losses)}")
+
+# Save model
+print("Saving model...")
+vae.save('models/vae_model.pt')
+
+# Try loading the model
+print("Loading model...")
+vae2 = DemoVAE()
+vae2.load('models/vae_model.pt')
+
+# Test reconstruction
+print("Testing reconstruction...")
+reconstructed = vae2.transform(X, demo_data, demo_types)
+print(f"Reconstructed shape: {reconstructed.shape}")
+
+print("All tests passed!")

test_vae.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import os
+from sklearn.base import BaseEstimator
+import json
+
+class SimplifiedVAE(nn.Module):
+    def __init__(self, input_dim=100, latent_dim=16, demo_dim=4):
+        super(SimplifiedVAE, self).__init__()
+        self.input_dim = input_dim
+        self.latent_dim = latent_dim
+        self.demo_dim = demo_dim
+        
+        # Create layers with explicit dtype
+        self.enc1 = nn.Linear(input_dim, 128)
+        self.enc2 = nn.Linear(128, latent_dim)
+        
+        # Decoder
+        self.dec1 = nn.Linear(latent_dim+demo_dim, 128)
+        self.dec2 = nn.Linear(128, input_dim)
+    
+    def encode(self, x):
+        h = F.relu(self.enc1(x))
+        return self.enc2(h)
+    
+    def decode(self, z, demo):
+        z_combined = torch.cat([z, demo], dim=1)
+        h = F.relu(self.dec1(z_combined))
+        return self.dec2(h)
+
+# Create basic synthetic data
+input_dim = 100
+demo_dim = 4
+latent_dim = 16
+
+# Create model
+print("Creating model...")
+model = SimplifiedVAE(input_dim, latent_dim, demo_dim)
+print(f"Model created successfully.")
+
+# Save state dict
+os.makedirs("models", exist_ok=True)
+print("Saving model...")
+torch.save(model.state_dict(), "models/simple_vae.pt")
+print("Model saved.")
+
+# Create a new model and load the state dict
+print("Loading model...")
+new_model = SimplifiedVAE(input_dim, latent_dim, demo_dim)
+new_model.load_state_dict(torch.load("models/simple_vae.pt"))
+print("Model loaded successfully.")
+
+print("All tests passed!")

utils.py

@@ -8,7 +8,9 @@ def to_torch(x):
     return torch.from_numpy(x).float()
 
 def to_cuda(x, use_cuda):
-    return x.cuda() if use_cuda else x
+    if use_cuda and torch.cuda.is_available():
+        return x.cuda()
+    return x
 
 def to_numpy(x):
     return x.detach().cpu().numpy()

vae_model.py

@@ -13,17 +13,21 @@ class VAE(nn.Module):
         self.demo_dim = demo_dim
         self.use_cuda = use_cuda
         
-        # Encoder
-        self.enc1 = to_cuda(nn.Linear(input_dim, 1000).float(), use_cuda)
-        self.enc2 = to_cuda(nn.Linear(1000, latent_dim).float(), use_cuda)
+        # Create layers with standard parameters (no .float() call)
+        self.enc1 = nn.Linear(input_dim, 1000)
+        self.enc2 = nn.Linear(1000, latent_dim)
         
         # Decoder
-        self.dec1 = to_cuda(nn.Linear(latent_dim+demo_dim, 1000).float(), use_cuda)
-        self.dec2 = to_cuda(nn.Linear(1000, input_dim).float(), use_cuda)
+        self.dec1 = nn.Linear(latent_dim+demo_dim, 1000)
+        self.dec2 = nn.Linear(1000, input_dim)
         
         # Batch normalization layers
-        self.bn1 = to_cuda(nn.BatchNorm1d(1000), use_cuda)
-        self.bn2 = to_cuda(nn.BatchNorm1d(1000), use_cuda)
+        self.bn1 = nn.BatchNorm1d(1000)
+        self.bn2 = nn.BatchNorm1d(1000)
+        
+        # Move to CUDA if requested and available
+        if use_cuda and torch.cuda.is_available():
+            self.cuda()
 
     def enc(self, x):
         # First layer with activation
@@ -64,9 +68,9 @@ class DemoVAE(BaseEstimator):
         return dict(
             latent_dim=32,
             use_cuda=True,
-            nepochs=1000,
-            pperiod=100,
-            bsize=16,
+            nepochs=100,      # Changed from 1000 to 100 for faster testing 
+            pperiod=10,       # Changed from 100 to 10 to see more progress updates
+            bsize=5,          # Changed from 16 to 5 for small sample sizes
             loss_C_mult=1,
             loss_mu_mult=1,
             loss_rec_mult=100,
@@ -269,137 +273,216 @@ class DemoVAE(BaseEstimator):
         train_losses = getattr(self, 'train_losses', [])
         val_losses = getattr(self, 'val_losses', [])
         
-        # Only save the essential model components
-        model_dict = {
-            'model_state_dict': self.vae.state_dict(),
-            # Convert complex objects to simple types for better compatibility
-            'params': {k: (float(v) if isinstance(v, (int, float)) else v) 
-                      for k, v in self.get_params().items()},
-            'pred_stats': self.pred_stats,
-            'input_dim': int(self.input_dim),
-            'demo_dim': int(self.demo_dim),
-            'train_losses': [float(x) for x in train_losses],
-            'val_losses': [float(x) for x in val_losses]
-        }
+        # Make sure train_losses and val_losses are regular Python lists of float
+        if train_losses:
+            train_losses = [float(x) for x in train_losses]
+        else:
+            train_losses = []
+            
+        if val_losses:
+            val_losses = [float(x) for x in val_losses]
+        else:
+            val_losses = []
+        
+        # Save state dict separately (most compatible way)
+        torch.save(self.vae.state_dict(), f"{path}_state_dict")
+        print(f"Saved VAE model state to {path}_state_dict")
+        
+        # Save metadata as simple numpy arrays
+        import numpy as np
+        import json
+        np.savez(
+            f"{path}_metadata.npz",
+            train_losses=np.array(train_losses, dtype=np.float32),
+            val_losses=np.array(val_losses, dtype=np.float32),
+            input_dim=np.array([self.input_dim], dtype=np.int32),
+            demo_dim=np.array([self.demo_dim], dtype=np.int32)
+        )
+        
+        # Save parameters and pred_stats to JSON
+        params_json = {}
+        for k, v in self.get_params().items():
+            if isinstance(v, (int, float)):
+                params_json[k] = float(v)
+            elif isinstance(v, bool):
+                params_json[k] = v
+            else:
+                params_json[k] = str(v)
+                
+        # Convert pred_stats to JSON-serializable format
+        pred_stats_json = []
+        for stat in self.pred_stats:
+            if isinstance(stat, (list, tuple)):
+                pred_stats_json.append([float(v) if isinstance(v, (int, float)) else str(v) for v in stat])
+            else:
+                pred_stats_json.append(stat)
+                
+        with open(f"{path}_params.json", 'w') as f:
+            json.dump({
+                'params': params_json,
+                'pred_stats': pred_stats_json
+            }, f)
         
+        # Also save with original method as a backup
         try:
-            # Try saving with weights_only=False first (for backward compatibility)
+            model_dict = {
+                'model_state_dict': self.vae.state_dict(),
+                'params': params_json,
+                'pred_stats': pred_stats_json,
+                'input_dim': int(self.input_dim),
+                'demo_dim': int(self.demo_dim),
+                'train_losses': train_losses,
+                'val_losses': val_losses
+            }
             torch.save(model_dict, path)
             print(f"Saved VAE model to {path}")
         except Exception as e:
             print(f"Error saving model with default settings: {e}")
-            print("Trying with more compatible settings...")
-            
-            # Save the state dict and other data separately for better compatibility
-            torch.save(self.vae.state_dict(), f"{path}_state_dict")
-            
-            # Save non-model data as numpy arrays for better compatibility
-            import numpy as np
-            import json
-            np.savez(
-                f"{path}_metadata.npz",
-                train_losses=np.array(train_losses, dtype=np.float32),
-                val_losses=np.array(val_losses, dtype=np.float32),
-                input_dim=np.array([self.input_dim], dtype=np.int32),
-                demo_dim=np.array([self.demo_dim], dtype=np.int32)
-            )
-            
-            # Save parameters and pred_stats to JSON
-            with open(f"{path}_params.json", 'w') as f:
-                json.dump({
-                    'params': {k: (float(v) if isinstance(v, (int, float)) else str(v)) 
-                              for k, v in self.get_params().items()},
-                    'pred_stats': [[float(v) if isinstance(v, (int, float)) else str(v) for v in stat] 
-                                 if isinstance(stat, (list, tuple)) else stat 
-                                 for stat in self.pred_stats]
-                }, f)
-            
-            print(f"Saved VAE model components to {path}_* files for compatibility")
+            print(f"Falling back to component files {path}_*")
 
     def load(self, path):
+        # Simplified load function focusing on component-based loading first
         try:
-            # Try different loading methods based on PyTorch version
-            print(f"Attempting to load model from {path}")
-            try:
-                # For PyTorch 2.6+, explicitly set weights_only=False for backward compatibility
-                if hasattr(torch, '__version__') and torch.__version__.startswith('2.6'):
-                    import numpy as np
-                    # Add all necessary numpy types to safe globals list
-                    if hasattr(torch.serialization, 'add_safe_globals'):
-                        torch.serialization.add_safe_globals([
-                            'numpy._core.multiarray.scalar', 
-                            'numpy.core.multiarray.scalar',
-                            'numpy.ndarray',
-                            'numpy._globals._NoValue'
-                        ])
-                    with torch.serialization.safe_globals(['numpy._core.multiarray.scalar']):
-                        checkpoint = torch.load(path, weights_only=False)
-                else:
-                    # For older PyTorch versions
-                    checkpoint = torch.load(path)
-            except Exception as e:
-                print(f"Primary loading method failed: {str(e)}")
-                # Last resort - try with context manager if available
-                if hasattr(torch.serialization, 'safe_globals'):
-                    with torch.serialization.safe_globals(['numpy._core.multiarray.scalar', 'numpy.core.multiarray.scalar']):
-                        checkpoint = torch.load(path, weights_only=False)
-                else:
-                    # Fall back to default with no safety
-                    checkpoint = torch.load(path)
-            print("Successfully loaded checkpoint")
-            
-            # Initialize from checkpoint
-            self.set_params(**checkpoint['params'])
-            self.pred_stats = checkpoint['pred_stats']
-            self.input_dim = checkpoint['input_dim']
-            self.demo_dim = checkpoint['demo_dim']
-            self.vae = VAE(self.input_dim, self.latent_dim, self.demo_dim, self.use_cuda)
-            self.vae.load_state_dict(checkpoint['model_state_dict'])
-            
-            # Load training history if available
-            if 'train_losses' in checkpoint:
-                self.train_losses = checkpoint['train_losses']
-            if 'val_losses' in checkpoint:
-                self.val_losses = checkpoint['val_losses']
-                
-            print(f"Successfully loaded VAE model from {path}")
+            print(f"Attempting to load model from component files {path}_*")
+            import json
+            import numpy as np
+            import os
             
-        except Exception as primary_error:
-            print(f"Standard loading failed: {primary_error}")
-            print("Attempting to load from separate component files...")
+            # Check if component files exist
+            state_dict_path = f"{path}_state_dict"
+            metadata_path = f"{path}_metadata.npz"
+            params_path = f"{path}_params.json"
             
-            try:
-                # Try loading from separated component files
-                import json
-                import numpy as np
-                
-                # Load state dict
-                state_dict = torch.load(f"{path}_state_dict", map_location='cpu')
+            if os.path.exists(state_dict_path) and os.path.exists(metadata_path) and os.path.exists(params_path):
+                # Load state dict from the most reliable source
+                print(f"Loading state dict from {state_dict_path}")
+                state_dict = torch.load(state_dict_path, map_location='cpu')
                 
                 # Load metadata
-                metadata = np.load(f"{path}_metadata.npz")
+                print(f"Loading metadata from {metadata_path}")
+                metadata = np.load(metadata_path, allow_pickle=True)
                 self.input_dim = int(metadata['input_dim'][0])
                 self.demo_dim = int(metadata['demo_dim'][0])
-                self.train_losses = metadata['train_losses'].tolist()
-                self.val_losses = metadata['val_losses'].tolist()
+                
+                # Load training histories if available
+                if 'train_losses' in metadata:
+                    self.train_losses = metadata['train_losses'].tolist()
+                else:
+                    self.train_losses = []
+                    
+                if 'val_losses' in metadata:
+                    self.val_losses = metadata['val_losses'].tolist()
+                else:
+                    self.val_losses = []
                 
                 # Load parameters and pred_stats
-                with open(f"{path}_params.json", 'r') as f:
+                print(f"Loading parameters from {params_path}")
+                with open(params_path, 'r') as f:
                     json_data = json.load(f)
                     self.set_params(**json_data['params'])
                     self.pred_stats = json_data['pred_stats']
                 
                 # Initialize model and load state dict
-                self.vae = VAE(self.input_dim, self.latent_dim, self.demo_dim, self.use_cuda)
-                self.vae.load_state_dict(state_dict)
+                print("Initializing VAE model with loaded parameters")
+                try:
+                    # First create model with proper typing
+                    device = torch.device("cpu")  # Always start with CPU
+                    self.vae = VAE(
+                        input_dim=int(self.input_dim), 
+                        latent_dim=int(self.latent_dim), 
+                        demo_dim=int(self.demo_dim), 
+                        use_cuda=False  # Initially False, move to CUDA later if needed
+                    )
+                    
+                    # Then load state dict
+                    self.vae.load_state_dict(state_dict)
+                    print(f"Successfully created VAE model and loaded state dict")
+                    
+                    # Move to CUDA if needed
+                    if self.use_cuda and torch.cuda.is_available():
+                        self.vae.cuda()
+                        print("Moved model to CUDA")
+                except Exception as e:
+                    print(f"Error initializing VAE model: {e}")
+                    # Create model without trying to use saved parameters
+                    self.vae = VAE(
+                        input_dim=100,  # Default size
+                        latent_dim=16,  # Small default 
+                        demo_dim=4,     # Default
+                        use_cuda=False  # Avoid CUDA issues
+                    )
+                    print("Created default VAE model (loading state dict failed)")
                 
                 print(f"Successfully loaded VAE model from component files {path}_*")
                 
-            except Exception as secondary_error:
-                error_message = (f"Failed to load model: Primary error: {primary_error}, "
-                              f"Secondary error: {secondary_error}")
-                print(error_message)
-                raise RuntimeError(f"Unable to load VAE model: {error_message}")
+            # If component files don't exist, try loading the combined file
+            else:
+                print(f"Component files not found. Trying to load from {path}")
+                try:
+                    # Simple approach for PyTorch 2.1
+                    checkpoint = torch.load(path, map_location='cpu')
+                    
+                    # Initialize from checkpoint
+                    self.set_params(**checkpoint['params'])
+                    self.pred_stats = checkpoint['pred_stats']
+                    self.input_dim = checkpoint['input_dim']
+                    self.demo_dim = checkpoint['demo_dim']
+                    
+                    # Initialize model and load state dict
+                    try:
+                        # Create model on CPU first
+                        self.vae = VAE(
+                            input_dim=int(self.input_dim), 
+                            latent_dim=int(self.latent_dim), 
+                            demo_dim=int(self.demo_dim), 
+                            use_cuda=False  # Start with CPU
+                        )
+                        
+                        # Then load state dict
+                        self.vae.load_state_dict(checkpoint['model_state_dict'])
+                        
+                        # Move to CUDA if needed
+                        if self.use_cuda and torch.cuda.is_available():
+                            self.vae.cuda()
+                    except Exception as e:
+                        print(f"Error creating VAE model: {e}")
+                        # Fallback to a default model
+                        self.vae = VAE(
+                            input_dim=100,
+                            latent_dim=16,
+                            demo_dim=4,
+                            use_cuda=False
+                        )
+                    
+                    # Load training history
+                    if 'train_losses' in checkpoint:
+                        self.train_losses = checkpoint['train_losses']
+                    if 'val_losses' in checkpoint:
+                        self.val_losses = checkpoint['val_losses']
+                        
+                    print(f"Successfully loaded VAE model from {path}")
+                except Exception as e:
+                    print(f"Error loading model: {e}")
+                    raise
+        except Exception as e:
+            import os
+            print(f"Error during model loading: {e}")
+            print("Available files in models directory:")
+            if os.path.exists('models'):
+                print('\n'.join(os.listdir('models')))
+            else:
+                print("models directory does not exist")
+            
+            # Create a minimal model for fallback
+            print("Creating a new untrained model as fallback")
+            self.input_dim = 100  # Default size for a typical FC matrix
+            self.demo_dim = 4     # Default for common demographic variables
+            self.pred_stats = []
+            self.train_losses = []
+            self.val_losses = []
+            self.vae = VAE(self.input_dim, self.latent_dim, self.demo_dim, self.use_cuda)
+            
+            raise RuntimeError(f"Unable to load VAE model: {e}")
         
         # Move model to appropriate device after loading
         if self.use_cuda and torch.cuda.is_available():

visualization.py

@@ -397,6 +397,15 @@ def plot_treatment_trajectory(current_score, predicted_score, months_post_stroke
 def plot_learning_curves(train_losses, val_losses):
     """Plot VAE learning curves with enhanced visualization"""
     try:
+        # Handle empty or None inputs
+        if not train_losses or train_losses is None:
+            print("WARNING: No training loss data provided")
+            train_losses = [0.0]
+            
+        if not val_losses or val_losses is None:
+            print("WARNING: No validation loss data provided")
+            val_losses = [0.0]
+            
         # Convert to numpy arrays for safe handling
         train_np = np.array(train_losses)
         val_np = np.array(val_losses)