Upload 13 files

README.md

@@ -1,12 +1,51 @@
----
-title: AphasiaPred
-emoji: 😻
-colorFrom: indigo
-colorTo: purple
-sdk: gradio
-sdk_version: 5.20.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Aphasia fMRI to FC Analysis using VAE
+
+This demo performs functional connectivity analysis on fMRI data using a Variational Autoencoder (VAE) approach. It's designed to work with aphasia patient data, analyzing brain connectivity patterns and their relationship to demographic variables.
+
+## About the Model
+
+This application implements a VAE model that:
+1. Takes functional connectivity (FC) matrices derived from fMRI data
+2. Learns a lower-dimensional latent representation of brain connectivity
+3. Conditions the generation process on demographic variables (age, sex, time post-stroke, WAB scores)
+4. Allows analysis of relationships between brain connectivity patterns and demographic variables
+
+## Dataset
+
+This demo uses the [SreekarB/OSFData](https://huggingface.co/datasets/SreekarB/OSFData) dataset from HuggingFace, which contains:
+
+- Functional connectivity matrices from fMRI data
+- Demographic information in `FC_graph_covariate_data.csv` including:
+  - Age at stroke
+  - Sex
+  - Months post-stroke
+  - WAB scores (aphasia severity)
+
+## How to Use
+
+1. **Data Source**: By default, it uses the HuggingFace dataset. You can change to a local directory if needed.
+2. **Model Parameters**:
+   - Latent Dimensions: Controls the size of the latent space (default: 32)
+   - Number of Epochs: Training iterations (default: 1000)
+   - Batch Size: Training batch size (default: 16)
+
+3. **Run the Analysis**: The model will:
+   - Load and process the data
+   - Train the VAE model
+   - Analyze relationships between latent variables and demographics
+   - Generate visualizations of original, reconstructed, and generated FC matrices
+
+## Outputs
+
+The application produces visualizations showing:
+- Original FC matrix
+- Reconstructed FC matrix
+- Generated FC matrix (based on specific demographic inputs)
+- Correlation plots between latent variables and demographic features
+
+## Technical Details
+
+- Framework: PyTorch
+- Interface: Gradio
+- Dataset: HuggingFace Datasets API
+- Analysis: Custom implementation of conditional VAE with demographic conditioning

analysis.py

@@ -0,0 +1,16 @@
+from scipy.stats import pearsonr
+
+def analyze_fc_patterns(latents, demographics):
+    results = {}
+    for demo_name, demo_values in demographics.items():
+        if demo_name != 'sex':  # For continuous variables
+            correlations = []
+            p_values = []
+            for latent_dim in range(latents.shape[1]):
+                r, p = pearsonr(latents[:, latent_dim], demo_values)
+                correlations.append(r)
+                p_values.append(p)
+            results[demo_name] = {'correlations': correlations, 'p_values': p_values}
+    
+    return results
+

app.py

@@ -0,0 +1,52 @@
+import gradio as gr
+from main import run_fc_analysis
+import os
+
+def gradio_fc_analysis(data_source, demographic_file, latent_dim, nepochs, bsize, use_hf_dataset):
+    fig = run_fc_analysis(
+        data_dir=data_source,
+        demographic_file=demographic_file,
+        latent_dim=latent_dim,
+        nepochs=nepochs,
+        bsize=bsize,
+        save_model=True,
+        use_hf_dataset=use_hf_dataset
+    )
+    return fig
+
+def create_interface():
+    iface = gr.Interface(
+        fn=gradio_fc_analysis,
+        inputs=[
+            gr.Textbox(label="Data Source (HF Dataset ID or Local Directory)", 
+                       default="SreekarB/OSFData"),
+            gr.Textbox(label="Demographic File", 
+                       default="FC_graph_covariate_data.csv"),
+            gr.Slider(minimum=8, maximum=64, step=8, 
+                      label="Latent Dimensions", default=32),
+            gr.Slider(minimum=100, maximum=5000, step=100, 
+                      label="Number of Epochs", default=1000),
+            gr.Slider(minimum=8, maximum=64, step=8, 
+                      label="Batch Size", default=16),
+            gr.Checkbox(label="Use HuggingFace Dataset", 
+                       default=True),
+        ],
+        outputs="plot",
+        title="Aphasia fMRI to FC Analysis using VAE",
+        description="""
+        Analysis pipeline: fMRI → FC matrices → VAE → Analysis
+        
+        This demo uses the SreekarB/OSFData dataset from HuggingFace by default.
+        The demographic file FC_graph_covariate_data.csv contains age_at_stroke, sex, months_post_stroke, and wab_score.
+        """,
+        examples=[
+            ["SreekarB/OSFData", "FC_graph_covariate_data.csv", 32, 500, 16, True],
+        ],
+        cache_examples=False,
+    )
+    return iface
+
+if __name__ == "__main__":
+    iface = create_interface()
+    iface.launch(share=True)
+

config.py

@@ -0,0 +1,24 @@
+# Model configuration
+MODEL_CONFIG = {
+    'latent_dim': 32,
+    'nepochs': 1000,
+    'bsize': 16,
+    'loss_rec_mult': 100,
+    'loss_decor_mult': 10,
+    'lr': 1e-4
+}
+
+# Preprocessing configuration
+PREPROCESS_CONFIG = {
+    't_r': 2.0,
+    'high_pass': 0.01,
+    'low_pass': 0.1,
+    'radius': 5
+}
+
+# Dataset configuration
+DATASET_CONFIG = {
+    'name': 'SreekarB/OSFData',
+    'split': 'train'
+}
+

data_preprocessing.py

@@ -0,0 +1,79 @@
+import numpy as np
+import pandas as pd
+from datasets import load_dataset
+from nilearn import input_data, connectome
+from nilearn.image import load_img
+import nibabel as nib
+
+def preprocess_fmri_to_fc(dataset_name, atlas_path=None):
+    dataset = load_dataset(dataset_name, split="train")
+    
+    # Load Power 264 atlas or specified atlas
+    if atlas_path is None:
+        # Use Power 264 coordinates to create spherical ROIs
+        from nilearn import datasets
+        power = datasets.fetch_coords_power_2011()
+        coords = np.vstack((power.rois['x'], power.rois['y'], power.rois['z'])).T
+        
+        masker = input_data.NiftiSpheresMasker(
+            coords, radius=5,
+            standardize=True,
+            memory='nilearn_cache', memory_level=1,
+            verbose=0,
+            detrend=True,
+            low_pass=0.1,
+            high_pass=0.01,
+            t_r=2.0  # Adjust TR according to your data
+        )
+    else:
+        masker = input_data.NiftiLabelsMasker(
+            labels_img=atlas_path,
+            standardize=True,
+            memory='nilearn_cache', memory_level=1,
+            verbose=0,
+            detrend=True,
+            low_pass=0.1,
+            high_pass=0.01,
+            t_r=2.0  # Adjust TR according to your data
+        )
+
+    # Load demographic data
+    demo_df = pd.DataFrame(dataset['demographics'])
+    
+    demo_data = [
+        demo_df['age_at_stroke'].values,
+        demo_df['sex'].values,
+        demo_df['months_post_stroke'].values,
+        demo_df['wab_score'].values
+    ]
+    
+    demo_types = ['continuous', 'categorical', 'continuous', 'continuous']
+    
+    # Process fMRI data and compute FC matrices
+    fc_matrices = []
+    for nii_file in dataset['nii_files']:
+        fmri_img = load_img(nii_file)
+        time_series = masker.fit_transform(fmri_img)
+        
+        correlation_measure = connectome.ConnectivityMeasure(
+            kind='correlation',
+            vectorize=False,
+            discard_diagonal=False
+        )
+        
+        fc_matrix = correlation_measure.fit_transform([time_series])[0]
+        
+        triu_indices = np.triu_indices_from(fc_matrix, k=1)
+        fc_triu = fc_matrix[triu_indices]
+        
+        fc_triu = np.arctanh(fc_triu)  # Fisher z-transform
+        
+        fc_matrices.append(fc_triu)
+    
+    X = np.array(fc_matrices)
+    
+    # Normalize the FC data
+    X = (X - np.mean(X, axis=0)) / np.std(X, axis=0)
+    
+    return X, demo_data, demo_types
+

main.py

@@ -0,0 +1,185 @@
+import os
+import numpy as np
+import torch
+from pathlib import Path
+import nibabel as nib
+from data_preprocessing import preprocess_fmri_to_fc
+from vae_model import train_fc_vae, DemoVAE
+from analysis import analyze_fc_patterns
+from visualization import visualize_fc_analysis
+from config import MODEL_CONFIG, DATASET_CONFIG
+import pandas as pd
+import io
+from typing import List, Dict, Union, Tuple, Any
+
+def load_data(data_dir="SreekarB/OSFData", demographic_file="FC_graph_covariate_data.csv", use_hf_dataset=True):
+    """
+    Load fMRI data and demographics from HuggingFace dataset or local files
+    """
+    if use_hf_dataset:
+        # Load from HuggingFace Datasets
+        from datasets import load_dataset
+        
+        print(f"Loading dataset from HuggingFace: {data_dir}")
+        dataset = load_dataset(data_dir)
+        
+        # Load demographics from the dataset
+        if demographic_file in dataset["train"].features:
+            demo_df = pd.DataFrame(dataset["train"][demographic_file])
+        else:
+            # Try to load from the dataset files
+            try:
+                demo_content = dataset["train"][demographic_file][0]
+                demo_df = pd.read_csv(io.StringIO(demo_content))
+            except Exception as e:
+                print(f"Error loading demographics from dataset: {e}")
+                # Download the CSV from the dataset repo
+                import huggingface_hub
+                csv_path = huggingface_hub.hf_hub_download(repo_id=data_dir, filename=demographic_file)
+                demo_df = pd.read_csv(csv_path)
+        
+        # Extract demographic data
+        demo_data = [
+            demo_df['age_at_stroke'].values if 'age_at_stroke' in demo_df.columns else np.array([]),
+            demo_df['sex'].values if 'sex' in demo_df.columns else np.array([]),
+            demo_df['months_post_stroke'].values if 'months_post_stroke' in demo_df.columns else np.array([]),
+            demo_df['wab_score'].values if 'wab_score' in demo_df.columns else np.array([])
+        ]
+        
+        # Get fMRI/FC files from dataset
+        nii_files = []
+        for f in dataset["train"].features:
+            if f.endswith(".nii.gz") or f.endswith(".nii"):
+                nii_files.append(f)
+        
+        if not nii_files:
+            print("No .nii/.nii.gz files found in dataset, checking for FC matrices")
+            # Try to find FC matrices directly
+            fc_matrices = []
+            for f in dataset["train"].features:
+                if f.startswith("fc_") or f.endswith("_fc"):
+                    fc_matrices.append(dataset["train"][f])
+            
+            if fc_matrices:
+                print(f"Found {len(fc_matrices)} FC matrices in dataset")
+                return fc_matrices, demo_data, demo_types
+    else:
+        # Original local file loading
+        # Load demographics
+        demo_df = pd.read_csv(demographic_file)
+        
+        demo_data = [
+            demo_df['age_at_stroke'].values,
+            demo_df['sex'].values,
+            demo_df['months_post_stroke'].values,
+            demo_df['wab_score'].values
+        ]
+        
+        # Load fMRI files
+        nii_files = sorted(list(Path(data_dir).glob('*.nii.gz')))
+    
+    demo_types = ['continuous', 'categorical', 'continuous', 'continuous']
+    return nii_files, demo_data, demo_types
+
+def run_fc_analysis(data_dir="SreekarB/OSFData", 
+                    demographic_file="FC_graph_covariate_data.csv", 
+                    latent_dim=32, 
+                    nepochs=1000, 
+                    bsize=16,
+                    save_model=True,
+                    use_hf_dataset=True):
+    
+    # Update MODEL_CONFIG with user-specified parameters
+    MODEL_CONFIG.update({
+        'latent_dim': latent_dim,
+        'nepochs': nepochs,
+        'bsize': bsize
+    })
+    
+    # Load data
+    print("Loading data...")
+    nii_files, demo_data, demo_types = load_data(data_dir, demographic_file, use_hf_dataset)
+    
+    # Add import for io module if it's missing
+    import io
+    
+    # Check if we got FC matrices directly
+    if isinstance(nii_files, list) and all(isinstance(item, np.ndarray) for item in nii_files):
+        print("Using pre-computed FC matrices...")
+        X = np.stack(nii_files)
+    else:
+        # Prepare data by converting fMRI to FC matrices
+        print("Converting fMRI data to FC matrices...")
+        X, demo_data, demo_types = preprocess_fmri_to_fc(nii_files, demo_data, demo_types)
+    
+    # Train VAE and get data
+    print("Training VAE...")
+    vae, X, demo_data, demo_types = train_fc_vae(X, demo_data, demo_types, MODEL_CONFIG)
+    
+    if save_model:
+        print("Saving model...")
+        os.makedirs('models', exist_ok=True)
+        torch.save(vae.state_dict(), 'models/vae_model.pth')
+    
+    # Get latent representations
+    print("Getting latent representations...")
+    latents = vae.get_latents(X)
+    
+    # Analyze results
+    print("Analyzing demographic relationships...")
+    demographics = {
+        'age_at_stroke': demo_data[0] if len(demo_data[0]) > 0 else np.zeros(len(X)),
+        'months_post_stroke': demo_data[2] if len(demo_data[2]) > 0 else np.zeros(len(X)),
+        'wab_score': demo_data[3] if len(demo_data[3]) > 0 else np.zeros(len(X))
+    }
+    analysis_results = analyze_fc_patterns(latents, demographics)
+    
+    # Generate new FC matrix
+    print("Generating new FC matrices...")
+    new_demographics = [
+        [60.0],  # age at stroke
+        ['M'],   # sex
+        [12.0],  # months post stroke
+        [80.0]   # wab score
+    ]
+    generated_fc = vae.transform(1, new_demographics, demo_types)
+    reconstructed_fc = vae.transform(X, demo_data, demo_types)
+    
+    # Visualize results
+    print("Creating visualizations...")
+    fig = visualize_fc_analysis(X[0], reconstructed_fc[0], generated_fc[0], analysis_results)
+    
+    return fig
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Run FC Analysis using VAE')
+    parser.add_argument('--data_dir', type=str, default='SreekarB/OSFData',
+                        help='HuggingFace dataset ID or directory containing fMRI data')
+    parser.add_argument('--demographic_file', type=str, default='FC_graph_covariate_data.csv',
+                        help='Path to demographic data CSV file')
+    parser.add_argument('--latent_dim', type=int, default=32,
+                        help='Dimension of latent space')
+    parser.add_argument('--nepochs', type=int, default=1000,
+                        help='Number of training epochs')
+    parser.add_argument('--bsize', type=int, default=16,
+                        help='Batch size for training')
+    parser.add_argument('--no_save', action='store_false',
+                        help='Do not save the model')
+    parser.add_argument('--use_local', action='store_true',
+                        help='Use local data instead of HuggingFace dataset')
+    
+    args = parser.parse_args()
+    
+    fig = run_fc_analysis(
+        data_dir=args.data_dir,
+        demographic_file=args.demographic_file,
+        latent_dim=args.latent_dim,
+        nepochs=args.nepochs,
+        bsize=args.bsize,
+        save_model=args.no_save,
+        use_hf_dataset=not args.use_local
+    )
+    fig.show()
+

requirements.txt

@@ -0,0 +1,12 @@
+torch>=1.9.0
+numpy>=1.19.2
+pandas>=1.2.4
+nilearn>=0.8.1
+nibabel>=3.2.1
+scikit-learn>=0.24.2
+matplotlib>=3.4.2
+gradio>=2.0.0
+datasets>=1.11.0
+huggingface_hub>=0.12.0
+transformers>=4.15.0
+

src/demovae/model.py

@@ -0,0 +1,221 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import random
+import numpy as np
+
+from sklearn.linear_model import Ridge
+from sklearn.linear_model import LogisticRegression
+
+def to_torch(x):
+    return torch.from_numpy(x).float()
+
+def to_cuda(x, use_cuda):
+    if use_cuda:
+        return x.cuda()
+    else:
+        return x
+
+def to_numpy(x):
+    return x.detach().cpu().numpy()
+
+class VAE(nn.Module):
+    def __init__(self, input_dim, latent_dim, demo_dim, use_cuda=True):
+        super(VAE, self).__init__()
+        self.input_dim = input_dim
+        self.latent_dim = latent_dim
+        self.demo_dim = demo_dim
+        self.use_cuda = use_cuda
+        self.enc1 = to_cuda(nn.Linear(input_dim, 1000).float(), use_cuda)
+        self.enc2 = to_cuda(nn.Linear(1000, latent_dim).float(), use_cuda)
+        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)
+
+    def enc(self, x):
+        x = F.relu(self.enc1(x))
+        z = self.enc2(x)
+        return z
+
+    def gen(self, n):
+        return to_cuda(torch.randn(n, self.latent_dim).float(), self.use_cuda)
+
+    def dec(self, z, demo):
+        z = to_cuda(torch.cat([z, demo], dim=1), self.use_cuda)
+        x = F.relu(self.dec1(z))
+        x = self.dec2(x)
+        #x = x.reshape(len(z), 264, 5)
+        #x = torch.einsum('nac,nbc->nab', x, x)
+        #a,b = np.triu_indices(264, 1)
+        #x = x[:,a,b]
+        return x
+
+def rmse(a, b, mean=torch.mean):
+    return mean((a-b)**2)**0.5
+
+def latent_loss(z, use_cuda=True):
+    C = z.T@z
+    mu = torch.mean(z, dim=0)
+    tgt1 = to_cuda(torch.eye(z.shape[-1]).float(), use_cuda)*len(z)
+    tgt2 = to_cuda(torch.zeros(z.shape[-1]).float(), use_cuda)
+    loss_C = rmse(C, tgt1)
+    loss_mu = rmse(mu, tgt2)
+    return loss_C, loss_mu, C, mu
+
+def decor_loss(z, demo, use_cuda=True):
+    ps = []
+    losses = []
+    for di in range(demo.shape[1]):
+        d = demo[:,di]
+        d = d - torch.mean(d)
+        p = torch.einsum('n,nz->z', d, z)
+        p = p/torch.std(d)
+        p = p/torch.einsum('nz,nz->z', z, z)
+        tgt = to_cuda(torch.zeros(z.shape[-1]).float(), use_cuda)
+        loss = rmse(p, tgt)
+        losses.append(loss)
+        ps.append(p)
+    losses = torch.stack(losses)
+    return losses, ps
+
+def pretty(x):
+    return f'{round(float(x), 4)}'
+
+def demo_to_torch(demo, demo_types, pred_stats, use_cuda):
+    demo_t = []
+    demo_idx = 0
+    for d,t,s in zip(demo, demo_types, pred_stats):
+        if t == 'continuous':
+            demo_t.append(to_cuda(to_torch(d), use_cuda))
+        elif t == 'categorical':
+            for dd in d:
+                if dd not in s:
+                    print(f'Model not trained with value {dd} for categorical demographic {demo_idx}')
+                    raise Exception('Bad demographic')
+            for ss in s:
+                idx = (d == ss).astype('bool')
+                zeros = torch.zeros(len(d))
+                zeros[idx] = 1
+                demo_t.append(to_cuda(zeros, use_cuda))
+        demo_idx += 1
+    demo_t = torch.stack(demo_t).permute(1,0)
+    return demo_t
+
+def train_vae(vae, x, demo, demo_types, nepochs, pperiod, bsize, loss_C_mult, loss_mu_mult, loss_rec_mult, loss_decor_mult, loss_pred_mult, lr, weight_decay, alpha, LR_C, ret_obj):
+    # Get linear predictors for demographics
+    pred_w = []
+    pred_i = []
+    # Pred stats are mean and std for continuous, and a list of all values for categorical
+    pred_stats = []
+    for i,d,t in zip(range(len(demo)), demo, demo_types):
+        print(f'Fitting auxilliary guidance model for demographic {i} {t}...', end='')
+        if t == 'continuous':
+            pred_stats.append([np.mean(d), np.std(d)])
+            reg = Ridge(alpha=alpha).fit(x, d)
+            reg_w = to_cuda(to_torch(reg.coef_), vae.use_cuda)
+            reg_i = reg.intercept_
+            pred_w.append(reg_w)
+            pred_i.append(reg_i)
+        elif t == 'categorical':
+            pred_stats.append(sorted(list(set(list(d)))))
+            reg = LogisticRegression(C=LR_C).fit(x, d)
+            # Binary
+            if len(reg.coef_) == 1:
+                reg_w = to_cuda(to_torch(reg.coef_[0]), vae.use_cuda)
+                reg_i = reg.intercept_[0]
+                pred_w.append(-reg_w)
+                pred_i.append(-reg_i)
+                pred_w.append(reg_w)
+                pred_i.append(reg_i)
+            # Categorical
+            else:
+                for i in range(len(reg.coef_)):
+                    reg_w = to_cuda(to_torch(reg.coef_[i]), vae.use_cuda)
+                    reg_i = reg.intercept_[i]
+                    pred_w.append(reg_w)
+                    pred_i.append(reg_i)
+        else:
+            print(f'demographic type "{t}" not "continuous" or "categorical"')
+            raise Exception('Bad demographic type')
+        print(' done')
+    ret_obj.pred_stats = pred_stats
+    # Convert input to pytorch
+    print('Converting input to pytorch')
+    x = to_cuda(to_torch(x), vae.use_cuda)
+    # Convert demographics to pytorch
+    print('Converting demographics to pytorch')
+    demo_t = demo_to_torch(demo, demo_types, pred_stats, vae.use_cuda)
+    # Training loop
+    print('Beginning VAE training')
+    ce = nn.CrossEntropyLoss()
+    optim = torch.optim.Adam(vae.parameters(), lr=lr, weight_decay=weight_decay)
+    for e in range(nepochs):
+        for bs in range(0,len(x),bsize):
+            xb = x[bs:(bs+bsize)]
+            db = demo_t[bs:(bs+bsize)]
+            optim.zero_grad()
+            # Reconstruct
+            z = vae.enc(xb)
+            y = vae.dec(z, db)
+            loss_C, loss_mu, _, _ = latent_loss(z, vae.use_cuda)
+            loss_decor, _ = decor_loss(z, db, vae.use_cuda)
+            loss_decor = sum(loss_decor)
+            loss_rec = rmse(xb, y)
+            # Sample demographics
+            demo_gen = []
+            for s,t in zip(pred_stats, demo_types):
+                if t == 'continuous':
+                    mu = s[0]
+                    std = s[1]
+                    dd = torch.randn(100).float()
+                    dd = dd*std+mu
+                    dd = to_cuda(dd, vae.use_cuda)
+                    demo_gen.append(dd)
+                elif t == 'categorical':
+                    idx = random.randint(0, len(s)-1)
+                    for i in range(len(s)):
+                        if idx == i:
+                            dd = torch.ones(100).float()
+                        else:
+                            dd = torch.zeros(100).float()
+                        dd = to_cuda(dd, vae.use_cuda)
+                        demo_gen.append(dd)
+            demo_gen = torch.stack(demo_gen).permute(1,0)
+            # Generate
+            z = vae.gen(100)
+            y = vae.dec(z, demo_gen)
+            # Regressor/classifier guidance loss
+            losses_pred = []
+            idcs = []
+            dg_idx = 0
+            for s,t in zip(pred_stats, demo_types):
+                if t == 'continuous':
+                    yy = y@pred_w[dg_idx]+pred_i[dg_idx]
+                    loss = rmse(demo_gen[:,dg_idx], yy)
+                    losses_pred.append(loss)
+                    idcs.append(float(demo_gen[0,dg_idx]))
+                    dg_idx += 1
+                elif t == 'categorical':
+                    loss = 0
+                    for i in range(len(s)):
+                        yy = y@pred_w[dg_idx]+pred_i[dg_idx]
+                        loss += ce(torch.stack([-yy, yy], dim=1), demo_gen[:,dg_idx].long())
+                        idcs.append(int(demo_gen[0,dg_idx]))
+                        dg_idx += 1
+                        losses_pred.append(loss)
+            total_loss = loss_C_mult*loss_C + loss_mu_mult*loss_mu + loss_rec_mult*loss_rec + loss_decor_mult*loss_decor + loss_pred_mult*sum(losses_pred)
+            total_loss.backward()
+            optim.step()
+            if e%pperiod == 0 or e == nepochs-1:
+                print(f'Epoch {e} ', end='')
+                print(f'ReconLoss {pretty(loss_rec)} ', end='')
+                print(f'CovarianceLoss {pretty(loss_C)} ', end='')
+                print(f'MeanLoss {pretty(loss_mu)} ', end='')
+                print(f'DecorLoss {pretty(loss_decor)} ', end='')
+                losses_pred = [pretty(loss) for loss in losses_pred]
+                print(f'GuidanceTargets {idcs} GuidanceLosses {losses_pred} ', end='')
+                print()
+    print('Training complete.')
+                
+

src/demovae/sklearn.py

@@ -0,0 +1,123 @@
+
+from demovae.model import VAE, train_vae, to_torch, to_cuda, to_numpy, demo_to_torch
+
+from sklearn.base import BaseEstimator
+
+# For saving
+import torch
+
+class DemoVAE(BaseEstimator):
+    def __init__(self, **params):
+        self.set_params(**params)
+
+    @staticmethod
+    def get_default_params():
+        return dict(latent_dim=60,      # Latent dimension
+                use_cuda=True,          # GPU acceleration
+                nepochs=3000,           # Training epochs
+                pperiod=100,            # Epochs between printing updates 
+                bsize=1000,             # Batch size
+                loss_C_mult=1,          # Covariance loss (KL div)
+                loss_mu_mult=1,         # Mean loss (KL div)
+                loss_rec_mult=100,      # Reconstruction loss
+                loss_decor_mult=10,     # Latent-demographic decorrelation loss
+                loss_pred_mult=0.001,   # Classifier/regressor guidance loss
+                alpha=100,              # Regularization for continuous guidance models
+                LR_C=100,               # Regularization for categorical guidance models
+                lr=1e-4,                # Learning rate
+                weight_decay=0,         # L2 regularization for VAE model
+                )
+
+    def get_params(self, **params):
+        return dict(latent_dim=self.latent_dim,
+                use_cuda=self.use_cuda,
+                nepochs=self.nepochs,
+                pperiod=self.pperiod,
+                bsize=self.bsize,
+                loss_C_mult=self.loss_C_mult,
+                loss_mu_mult=self.loss_mu_mult,
+                loss_rec_mult=self.loss_rec_mult,
+                loss_decor_mult=self.loss_decor_mult,
+                loss_pred_mult=self.loss_pred_mult,
+                alpha=self.alpha,
+                LR_C=self.LR_C,
+                lr=self.lr,
+                weight_decay=self.weight_decay,
+                )
+
+    def set_params(self, **params):
+        dft = DemoVAE.get_default_params()
+        for key in dft:
+            if key in params:
+                setattr(self, key, params[key])
+            else:
+                setattr(self, key, dft[key])
+        return self
+
+    def fit(self, x, demo, demo_types, **kwargs):
+        # Get demo_dim
+        demo_dim = 0
+        for d,t in zip(demo, demo_types):
+            if t == 'continuous':
+                demo_dim += 1
+            elif t == 'categorical':
+                ll = len(set(list(d)))
+                if ll == 1:
+                    print('Only one type of category for categorical variable')
+                    raise Exception('Bad categorical')
+                demo_dim += ll
+            else:
+                print(f'demographic type "{t}" not "continuous" or "categorical"')
+                raise Exception('Bad demographic type')
+        # Save parameters
+        self.input_dim = x.shape[1]
+        self.demo_dim = demo_dim
+        # Create model
+        self.vae = VAE(x.shape[1], self.latent_dim, demo_dim, self.use_cuda)
+        # Train model
+        train_vae(self.vae, x, demo, demo_types, 
+                self.nepochs, self.pperiod, self.bsize, 
+                self.loss_C_mult, self.loss_mu_mult, self.loss_rec_mult, self.loss_decor_mult, self.loss_pred_mult,
+                self.lr, self.weight_decay, self.alpha, self.LR_C, 
+                self)
+        return self
+
+    def transform(self, x, demo, demo_types, **kwargs):
+        if isinstance(x, int):
+            # Generate
+            z = self.vae.gen(x)
+        else:
+            # Get latents for real data
+            z = self.vae.enc(to_cuda(to_torch(x), self.vae.use_cuda))
+        demo_t = demo_to_torch(demo, demo_types, self.pred_stats, self.vae.use_cuda)
+        y = self.vae.dec(z, demo_t)
+        return to_numpy(y)
+
+    def fit_transform(self, x, demo, demo_types, **kwargs):
+        self.fit(x, demo, demo_types)
+        return self.transform(x, demo, demo_types)
+
+    def get_latents(self, x):
+        z = self.vae.enc(to_cuda(to_torch(x), self.vae.use_cuda))
+        return to_numpy(z)
+
+    def save(self, path):
+        params = self.get_params()
+        dct = dict(pred_stats=self.pred_stats, 
+                   params=params, 
+                   input_dim=self.input_dim,
+                   demo_dim=self.demo_dim,
+                   model_state_dict=self.vae.state_dict())
+        torch.save(dct, path)
+
+    def load(self, path):
+        dct = torch.load(path)
+        self.pred_stats = dct['pred_stats']
+        self.set_params(**dct['params'])
+        self.vae = VAE(dct['input_dim'], 
+                       dct['params']['latent_dim'], 
+                       dct['demo_dim'], 
+                       dct['params']['use_cuda'])
+        self.vae.load_state_dict(dct['model_state_dict'])
+
+

utils.py

@@ -0,0 +1,186 @@
+import torch
+import numpy as np
+from sklearn.linear_model import Ridge, LogisticRegression
+
+def to_torch(x):
+    return torch.from_numpy(x).float()
+
+def to_cuda(x, use_cuda):
+    return x.cuda() if use_cuda else x
+
+def to_numpy(x):
+    return x.detach().cpu().numpy()
+
+def fc_matrix_from_triu(triu_values, n_rois=264):
+    fc_matrix = np.zeros((n_rois, n_rois))
+    triu_indices = np.triu_indices(n_rois, k=1)
+    triu_values = np.tanh(triu_values)
+    fc_matrix[triu_indices] = triu_values
+    fc_matrix = fc_matrix + fc_matrix.T
+    np.fill_diagonal(fc_matrix, 1)
+    return fc_matrix
+
+def rmse(a, b, mean=torch.mean):
+    return mean((a-b)**2)**0.5
+
+def latent_loss(z, use_cuda=True):
+    C = z.T@z
+    mu = torch.mean(z, dim=0)
+    tgt1 = to_cuda(torch.eye(z.shape[-1]).float(), use_cuda)*len(z)
+    tgt2 = to_cuda(torch.zeros(z.shape[-1]).float(), use_cuda)
+    loss_C = rmse(C, tgt1)
+    loss_mu = rmse(mu, tgt2)
+    return loss_C, loss_mu, C, mu
+
+def decor_loss(z, demo, use_cuda=True):
+    ps = []
+    losses = []
+    for di in range(demo.shape[1]):
+        d = demo[:,di]
+        d = d - torch.mean(d)
+        p = torch.einsum('n,nz->z', d, z)
+        p = p/torch.std(d)
+        p = p/torch.einsum('nz,nz->z', z, z)
+        tgt = to_cuda(torch.zeros(z.shape[-1]).float(), use_cuda)
+        loss = rmse(p, tgt)
+        losses.append(loss)
+        ps.append(p)
+    losses = torch.stack(losses)
+    return losses, ps
+def demo_to_torch(demo, demo_types, pred_stats, use_cuda):
+    demo_t = []
+    demo_idx = 0
+    for d, t, s in zip(demo, demo_types, pred_stats):
+        if t == 'continuous':
+            demo_t.append(to_cuda(to_torch(d), use_cuda))
+        elif t == 'categorical':
+            for dd in d:
+                if dd not in s:
+                    print(f'Model not trained with value {dd} for categorical demographic {demo_idx}')
+                    raise Exception('Bad demographic')
+            for ss in s:
+                idx = (d == ss).astype('bool')
+                zeros = torch.zeros(len(d))
+                zeros[idx] = 1
+                demo_t.append(to_cuda(zeros, use_cuda))
+        demo_idx += 1
+    demo_t = torch.stack(demo_t).permute(1,0)
+    return demo_t
+
+def train_vae(vae, x, demo, demo_types, nepochs, pperiod, bsize, 
+              loss_C_mult, loss_mu_mult, loss_rec_mult, loss_decor_mult, 
+              loss_pred_mult, lr, weight_decay, alpha, LR_C, ret_obj):
+    # Get linear predictors for demographics
+    pred_w = []
+    pred_i = []
+    pred_stats = []
+    
+    for i, d, t in zip(range(len(demo)), demo, demo_types):
+        print(f'Fitting auxiliary guidance model for demographic {i} {t}...', end='')
+        if t == 'continuous':
+            pred_stats.append([np.mean(d), np.std(d)])
+            reg = Ridge(alpha=alpha).fit(x, d)
+            reg_w = to_cuda(to_torch(reg.coef_), vae.use_cuda)
+            reg_i = reg.intercept_
+            pred_w.append(reg_w)
+            pred_i.append(reg_i)
+        elif t == 'categorical':
+            pred_stats.append(sorted(list(set(list(d)))))
+            reg = LogisticRegression(C=LR_C).fit(x, d)
+            if len(reg.coef_) == 1:
+                reg_w = to_cuda(to_torch(reg.coef_[0]), vae.use_cuda)
+                reg_i = reg.intercept_[0]
+                pred_w.append(-reg_w)
+                pred_i.append(-reg_i)
+                pred_w.append(reg_w)
+                pred_i.append(reg_i)
+            else:
+                for i in range(len(reg.coef_)):
+                    reg_w = to_cuda(to_torch(reg.coef_[i]), vae.use_cuda)
+                    reg_i = reg.intercept_[i]
+                    pred_w.append(reg_w)
+                    pred_i.append(reg_i)
+        print(' done')
+    
+    ret_obj.pred_stats = pred_stats
+    
+    # Convert input to pytorch
+    x = to_cuda(to_torch(x), vae.use_cuda)
+    
+    # Convert demographics to pytorch
+    demo_t = demo_to_torch(demo, demo_types, pred_stats, vae.use_cuda)
+    
+    # Training loop
+    ce = torch.nn.CrossEntropyLoss()
+    optim = torch.optim.Adam(vae.parameters(), lr=lr, weight_decay=weight_decay)
+    
+    for e in range(nepochs):
+        for bs in range(0, len(x), bsize):
+            xb = x[bs:(bs+bsize)]
+            db = demo_t[bs:(bs+bsize)]
+            optim.zero_grad()
+            
+            # Reconstruct
+            z = vae.enc(xb)
+            y = vae.dec(z, db)
+            loss_C, loss_mu, _, _ = latent_loss(z, vae.use_cuda)
+            loss_decor, _ = decor_loss(z, db, vae.use_cuda)
+            loss_decor = sum(loss_decor)
+            loss_rec = rmse(xb, y)
+            
+            # Sample demographics
+            demo_gen = []
+            for s, t in zip(pred_stats, demo_types):
+                if t == 'continuous':
+                    mu, std = s
+                    dd = torch.randn(100).float()
+                    dd = dd*std+mu
+                    dd = to_cuda(dd, vae.use_cuda)
+                    demo_gen.append(dd)
+                elif t == 'categorical':
+                    idx = np.random.randint(0, len(s))
+                    for i in range(len(s)):
+                        dd = torch.ones(100).float() if idx == i else torch.zeros(100).float()
+                        dd = to_cuda(dd, vae.use_cuda)
+                        demo_gen.append(dd)
+            
+            demo_gen = torch.stack(demo_gen).permute(1,0)
+            
+            # Generate
+            z = vae.gen(100)
+            y = vae.dec(z, demo_gen)
+            
+            # Regressor/classifier guidance loss
+            losses_pred = []
+            idcs = []
+            dg_idx = 0
+            
+            for s, t in zip(pred_stats, demo_types):
+                if t == 'continuous':
+                    yy = y@pred_w[dg_idx]+pred_i[dg_idx]
+                    loss = rmse(demo_gen[:,dg_idx], yy)
+                    losses_pred.append(loss)
+                    idcs.append(float(demo_gen[0,dg_idx]))
+                    dg_idx += 1
+                elif t == 'categorical':
+                    loss = 0
+                    for i in range(len(s)):
+                        yy = y@pred_w[dg_idx]+pred_i[dg_idx]
+                        loss += ce(torch.stack([-yy, yy], dim=1), demo_gen[:,dg_idx].long())
+                        idcs.append(int(demo_gen[0,dg_idx]))
+                        dg_idx += 1
+                        losses_pred.append(loss)
+            
+            total_loss = (loss_C_mult*loss_C + loss_mu_mult*loss_mu + 
+                         loss_rec_mult*loss_rec + loss_decor_mult*loss_decor + 
+                         loss_pred_mult*sum(losses_pred))
+            
+            total_loss.backward()
+            optim.step()
+            
+            if e%pperiod == 0 or e == nepochs-1:
+                print(f'Epoch {e} ReconLoss {loss_rec:.4f} CovarianceLoss {loss_C:.4f} '
+                      f'MeanLoss {loss_mu:.4f} DecorLoss {loss_decor:.4f}')
+                print(f'GuidanceTargets {idcs}')
+                print(f'GuidanceLosses {[f"{loss:.4f}" for loss in losses_pred]}')
+

vae_model.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from utils import to_torch, to_cuda, to_numpy, demo_to_torch
+from sklearn.base import BaseEstimator
+
+class VAE(nn.Module):
+    def __init__(self, input_dim, latent_dim, demo_dim, use_cuda=True):
+        super(VAE, self).__init__()
+        self.input_dim = input_dim
+        self.latent_dim = latent_dim
+        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)
+        
+        # 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)
+        
+        # Batch normalization layers
+        self.bn1 = to_cuda(nn.BatchNorm1d(1000), use_cuda)
+        self.bn2 = to_cuda(nn.BatchNorm1d(1000), use_cuda)
+
+    def enc(self, x):
+        x = self.bn1(F.relu(self.enc1(x)))
+        z = self.enc2(x)
+        return z
+
+    def gen(self, n):
+        return to_cuda(torch.randn(n, self.latent_dim).float(), self.use_cuda)
+
+    def dec(self, z, demo):
+        z = to_cuda(torch.cat([z, demo], dim=1), self.use_cuda)
+        x = self.bn2(F.relu(self.dec1(z)))
+        x = self.dec2(x)
+        return x
+
+class DemoVAE(BaseEstimator):
+    def __init__(self, **params):
+        self.set_params(**params)
+
+    @staticmethod
+    def get_default_params():
+        return dict(
+            latent_dim=32,
+            use_cuda=True,
+            nepochs=1000,
+            pperiod=100,
+            bsize=16,
+            loss_C_mult=1,
+            loss_mu_mult=1,
+            loss_rec_mult=100,
+            loss_decor_mult=10,
+            loss_pred_mult=0.001,
+            alpha=100,
+            LR_C=100,
+            lr=1e-4,
+            weight_decay=0
+        )
+
+    def get_params(self, deep=True):
+        return {k: getattr(self, k) for k in self.get_default_params().keys()}
+
+    def set_params(self, **params):
+        for k, v in self.get_default_params().items():
+            setattr(self, k, params.get(k, v))
+        return self
+
+    def fit(self, x, demo, demo_types):
+        from utils import train_vae
+        
+        # Calculate demo_dim
+        demo_dim = 0
+        for d, t in zip(demo, demo_types):
+            if t == 'continuous':
+                demo_dim += 1
+            elif t == 'categorical':
+                demo_dim += len(set(d))
+            else:
+                raise ValueError(f'Demographic type "{t}" not supported')
+        
+        # Initialize VAE
+        self.input_dim = x.shape[1]
+        self.demo_dim = demo_dim
+        self.vae = VAE(self.input_dim, self.latent_dim, demo_dim, self.use_cuda)
+        
+        # Train VAE
+        train_vae(
+            self.vae, x, demo, demo_types,
+            self.nepochs, self.pperiod, self.bsize,
+            self.loss_C_mult, self.loss_mu_mult, self.loss_rec_mult,
+            self.loss_decor_mult, self.loss_pred_mult,
+            self.lr, self.weight_decay, self.alpha, self.LR_C,
+            self
+        )
+        return self
+
+    def transform(self, x, demo, demo_types):
+        if isinstance(x, int):
+            z = self.vae.gen(x)
+        else:
+            z = self.vae.enc(to_cuda(to_torch(x), self.vae.use_cuda))
+        demo_t = demo_to_torch(demo, demo_types, self.pred_stats, self.vae.use_cuda)
+        y = self.vae.dec(z, demo_t)
+        return to_numpy(y)
+
+    def get_latents(self, x):
+        z = self.vae.enc(to_cuda(to_torch(x), self.vae.use_cuda))
+        return to_numpy(z)
+
+    def save(self, path):
+        torch.save({
+            'model_state_dict': self.vae.state_dict(),
+            'params': self.get_params(),
+            'pred_stats': self.pred_stats,
+            'input_dim': self.input_dim,
+            'demo_dim': self.demo_dim
+        }, path)
+
+    def load(self, path):
+        checkpoint = torch.load(path)
+        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'])
+
+def train_fc_vae(X, demo_data, demo_types, model_config):
+    n_rois = 264
+    input_dim = (n_rois * (n_rois - 1)) // 2
+    
+    vae = DemoVAE(
+        latent_dim=model_config['latent_dim'],
+        nepochs=model_config['nepochs'],
+        bsize=model_config['bsize'],
+        loss_rec_mult=model_config['loss_rec_mult'],
+        loss_decor_mult=model_config['loss_decor_mult'],
+        lr=model_config['lr'],
+        use_cuda=torch.cuda.is_available()
+    )
+    
+    vae.fit(X, demo_data, demo_types)
+    
+    return vae, X, demo_data, demo_types
+

visualization.py

@@ -0,0 +1,44 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from utils import fc_matrix_from_triu
+
+def visualize_fc_analysis(original_triu, reconstructed_triu, generated_triu, analysis_results=None):
+    fig = plt.figure(figsize=(15, 10))
+    gs = plt.GridSpec(2, 3)
+    
+    ax1 = fig.add_subplot(gs[0, 0])
+    ax2 = fig.add_subplot(gs[0, 1])
+    ax3 = fig.add_subplot(gs[0, 2])
+    
+    original = fc_matrix_from_triu(original_triu)
+    reconstructed = fc_matrix_from_triu(reconstructed_triu)
+    generated = fc_matrix_from_triu(generated_triu)
+    
+    im1 = ax1.imshow(original, cmap='RdBu_r', vmin=-1, vmax=1)
+    ax1.set_title('Original FC')
+    
+    im2 = ax2.imshow(reconstructed, cmap='RdBu_r', vmin=-1, vmax=1)
+    ax2.set_title('Reconstructed FC')
+    
+    im3 = ax3.imshow(generated, cmap='RdBu_r', vmin=-1, vmax=1)
+    ax3.set_title('Generated FC')
+    
+    plt.colorbar(im1, ax=ax1)
+    plt.colorbar(im2, ax=ax2)
+    plt.colorbar(im3, ax=ax3)
+    
+    if analysis_results is not None:
+        ax4 = fig.add_subplot(gs[1, :])
+        for demo_name, results in analysis_results.items():
+            significant_dims = np.where(np.array(results['p_values']) < 0.05)[0]
+            correlations = np.array(results['correlations'])
+            ax4.plot(correlations, label=f'{demo_name} (sig. dims: {len(significant_dims)})')
+        
+        ax4.set_xlabel('Latent Dimension')
+        ax4.set_ylabel('Correlation Strength')
+        ax4.set_title('Demographic Correlations with Latent Dimensions')
+        ax4.legend()
+    
+    plt.tight_layout()
+    return fig
+