Initial commit: GraPHFormer codebase

.gitignore

@@ -0,0 +1,22 @@
+# Logs, results, and data
+work_dir/
+data/
+
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+*.egg
+.eggs/
+dist/
+build/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Environment
+.env
+*.log
+
+# OS
+.DS_Store

finetune.py

@@ -0,0 +1,554 @@
+"""
+GraPHFormer Fine-tuning Script
+
+Fine-tune pretrained CLIP-style models for classification.
+Supports: image_only, tree_only, multimodal modes.
+
+Usage:
+    python finetune.py --exp_name my_finetune --pretrained_checkpoint path/to/checkpoint.pth
+"""
+
+import argparse
+import datetime
+import time
+import os
+import json
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.backends.cudnn as cudnn
+from torch.utils.data import DataLoader
+from sklearn.neighbors import KNeighborsClassifier
+
+from graphformer.models import CLIPModel, FineTuneModel
+from graphformer.augmentations import (
+    Compose,
+    RandomScaleCoords, RandomRotate, RandomJitter, RandomShift,
+    RandomFlip, RandomMaskFeats, RandomJitterLength, RandomElasticate,
+    RandomDropSubTrees, RandomSkipParentNode, RandomSwapSiblingSubTrees,
+    CombinedPersistenceAugmentation,
+)
+from graphformer.utils import save_checkpoint, get_root_logger, set_seed
+from graphformer.data import NeuronTreeDataset, get_collate_fn, LABEL_DICT
+
+
+def evaluate_accuracy(model, data_loader, device):
+    """Evaluate classification accuracy"""
+    model.eval()
+    correct = 0
+    total = 0
+
+    with torch.no_grad():
+        for batch in data_loader:
+            batch = batch.to(device)
+            _, logits = model(batch)
+
+            pred = logits.argmax(dim=1)
+            labels = batch.label.cuda() if not batch.label.is_cuda else batch.label
+
+            correct += (pred == labels).sum().item()
+            total += labels.size(0)
+
+    accuracy = correct / total * 100
+    model.train()
+    return accuracy
+
+
+def extract_features(model, data_loader, device):
+    """Extract features for KNN evaluation"""
+    model.eval()
+    features_list = []
+    labels_list = []
+
+    with torch.no_grad():
+        for batch in data_loader:
+            batch = batch.to(device)
+            _, _, features = model(batch, return_features=True)
+
+            features_list.append(features)
+            labels = batch.label.cuda() if not batch.label.is_cuda else batch.label
+            labels_list.append(labels)
+
+    features = torch.cat(features_list, dim=0)
+    labels = torch.cat(labels_list, dim=0)
+
+    model.train()
+    return features, labels
+
+
+def evaluate_knn(model, train_loader, test_loader, device, knn_k=20):
+    """KNN evaluation using sklearn"""
+    x_train, y_train = extract_features(model, train_loader, device)
+    x_test, y_test = extract_features(model, test_loader, device)
+
+    neigh = KNeighborsClassifier(n_neighbors=knn_k)
+    neigh.fit(x_train.cpu().numpy(), y_train.cpu().numpy())
+
+    score = neigh.score(x_test.cpu().numpy(), y_test.cpu().numpy())
+
+    return score * 100
+
+
+def mixup_data(features, labels, alpha=1.0):
+    """Apply mixup augmentation"""
+    if alpha > 0:
+        lam = np.random.beta(alpha, alpha)
+    else:
+        lam = 1
+
+    batch_size = features.size(0)
+    index = torch.randperm(batch_size).to(features.device)
+
+    mixed_features = lam * features + (1 - lam) * features[index, :]
+    labels_a, labels_b = labels, labels[index]
+
+    return mixed_features, labels_a, labels_b, lam
+
+
+def mixup_criterion(criterion, pred, labels_a, labels_b, lam):
+    """Mixup loss"""
+    return lam * criterion(pred, labels_a) + (1 - lam) * criterion(pred, labels_b)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("GraPHFormer Fine-tuning")
+
+    # Basic
+    parser.add_argument("--work_dir", type=str, default="./work_dir")
+    parser.add_argument("--exp_name", type=str, required=True)
+    parser.add_argument("--dataset", type=str, default="bil_6_classes")
+    parser.add_argument("--data_dir", type=str, default="data/raw/bil")
+    parser.add_argument("--seed", type=int, default=42)
+
+    # Pretrained checkpoint
+    parser.add_argument("--pretrained_checkpoint", type=str, default=None)
+
+    # Fine-tuning mode
+    parser.add_argument("--mode", type=str, default="multimodal",
+                        choices=["image_only", "tree_only", "multimodal"])
+    parser.add_argument("--freeze_encoders", action="store_true", default=False)
+    parser.add_argument("--freeze_image_only", action="store_true", default=False)
+    parser.add_argument("--linear_probe_epochs", type=int, default=0)
+    parser.add_argument("--use_projection", action="store_true", default=False)
+    parser.add_argument("--fusion_mode", type=str, default="concat",
+                        choices=["concat", "add", "cross_attention", "bi_attention", "gated", "cmf", "mhcma"])
+
+    # Tree Model
+    parser.add_argument("--tree_model", type=str, default="double",
+                        choices=["ori", "v2", "double"])
+    parser.add_argument("--child_mode", type=str, default="sum")
+    parser.add_argument("--input_features", nargs="+", type=int,
+                        default=[2, 3, 4, 12, 13])
+    parser.add_argument("--h_size", type=int, default=256)
+    parser.add_argument("--bn", action="store_true", default=False)
+
+    # Image Model
+    parser.add_argument("--image_encoder", type=str, default="resnet18")
+    parser.add_argument("--image_size", type=int, default=256)
+    parser.add_argument("--freeze_image_backbone", action="store_true", default=False)
+
+    # CLIP settings
+    parser.add_argument("--embed_dim", type=int, default=128)
+    parser.add_argument("--temperature", type=float, default=0.07)
+    parser.add_argument("--loss_type", type=str, default="clip")
+
+    # Training
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--epochs", default=50, type=int)
+    parser.add_argument("--lr", default=1e-4, type=float)
+    parser.add_argument("--wd", default=0.01, type=float)
+    parser.add_argument("--warmup_epochs", type=int, default=5)
+    parser.add_argument("--start_epoch", type=int, default=0)
+    parser.add_argument("--save_freq", type=int, default=10)
+    parser.add_argument("--val_freq", type=int, default=1)
+    parser.add_argument("--gpu", default=0, type=int)
+
+    # Regularization
+    parser.add_argument("--label_smoothing", type=float, default=0.0)
+    parser.add_argument("--dropout", type=float, default=0.5)
+    parser.add_argument("--mixup_alpha", type=float, default=0.0)
+    parser.add_argument("--early_stopping_patience", type=int, default=0)
+
+    # ArcFace loss
+    parser.add_argument("--use_arcface", action="store_true", default=False)
+    parser.add_argument("--arcface_s", type=float, default=30.0)
+    parser.add_argument("--arcface_m", type=float, default=0.50)
+
+    # Augmentation
+    parser.add_argument("--aug_scale_coords", action="store_true", default=False)
+    parser.add_argument("--aug_rotate", action="store_true", default=False)
+    parser.add_argument("--aug_jitter_coords", action="store_true", default=False)
+    parser.add_argument("--aug_shift_coords", action="store_true", default=False)
+    parser.add_argument("--aug_flip", action="store_true", default=False)
+    parser.add_argument("--aug_mask_feats", action="store_true", default=False)
+    parser.add_argument("--aug_jitter_length", action="store_true", default=False)
+    parser.add_argument("--aug_elasticate", action="store_true", default=False)
+    parser.add_argument("--aug_drop_tree", action="store_true", default=False)
+    parser.add_argument("--aug_skip_parent_node", action="store_true", default=False)
+    parser.add_argument("--aug_swap_sibling_subtrees", action="store_true", default=False)
+
+    # Persistence augmentation
+    parser.add_argument("--use_persistence_aug", action="store_true", default=False)
+    parser.add_argument("--pers_translation_scale", type=float, default=0.05)
+    parser.add_argument("--pers_noise_scale", type=float, default=0.02)
+    parser.add_argument("--pers_sigma_min", type=float, default=12.0)
+    parser.add_argument("--pers_sigma_max", type=float, default=20.0)
+    parser.add_argument("--sigma_px", type=float, default=16.0)
+
+    # Evaluation
+    parser.add_argument("--eval_mode", type=str, default="accuracy",
+                        choices=["accuracy", "knn"])
+    parser.add_argument("--knn_k", type=int, default=20)
+
+    parser.add_argument("--cache_images", action="store_true", default=True)
+    parser.add_argument("--debug", action="store_true", default=False)
+
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    if args.linear_probe_epochs > 0:
+        args.freeze_encoders = True
+        
+    # Setup work directory
+    args.work_dir = f"{args.work_dir}/{args.exp_name}"
+    if not os.path.exists(args.work_dir):
+        os.makedirs(args.work_dir)
+
+    # Logger
+    timestamp = time.strftime("%Y%m%d_%H%M%S", time.localtime())
+    if args.debug:
+        log_file = None
+        args.save_freq = 10000
+        args.val_freq = 1
+    else:
+        log_file = f"{args.work_dir}/finetune_{timestamp}.log"
+    logger = get_root_logger(log_file=log_file, log_level="INFO")
+
+    logger.info("=" * 60)
+    logger.info("GraPHFormer FINE-TUNING")
+    logger.info(f"Mode: {args.mode}")
+    if args.mode == "multimodal":
+        logger.info(f"Fusion Mode: {args.fusion_mode}")
+    logger.info(f"Freeze Encoders: {args.freeze_encoders}")
+    logger.info(f"Pretrained Checkpoint: {args.pretrained_checkpoint}")
+    logger.info(f"Dataset: {args.dataset}")
+    logger.info("=" * 60)
+    logger.info(json.dumps(vars(args), indent=4, sort_keys=True))
+
+    device = torch.device("cuda")
+
+    # Load pretrained model or create from scratch
+    if args.pretrained_checkpoint is not None:
+        logger.info("=> Loading pretrained model...")
+        if not os.path.isfile(args.pretrained_checkpoint):
+            raise FileNotFoundError(f"Checkpoint not found: {args.pretrained_checkpoint}")
+
+        checkpoint = torch.load(args.pretrained_checkpoint, map_location=f"cuda:{args.gpu}")
+        state_dict = checkpoint["state_dict"]
+
+        # Auto-detect image encoder
+        if args.mode in ['image_only', 'multimodal']:
+            if "image_encoder.encoder.backbone.cls_token" in state_dict:
+                args.image_encoder = "dinov2_vits14"
+                logger.info(f"=> Detected DINOv2 image encoder")
+
+        pretrained_model = CLIPModel(args).to(device)
+        missing_keys, unexpected_keys = pretrained_model.load_state_dict(state_dict, strict=False)
+
+        if missing_keys:
+            logger.warning(f"=> Missing keys: {len(missing_keys)}")
+        if unexpected_keys:
+            logger.warning(f"=> Unexpected keys: {len(unexpected_keys)}")
+
+        logger.info(f"=> Loaded checkpoint from epoch {checkpoint.get('epoch', 'unknown')}")
+    else:
+        logger.info("=> Training from scratch")
+        pretrained_model = CLIPModel(args).to(device)
+
+    # Setup augmentations
+    aug_switchs = [
+        False,
+        args.aug_scale_coords,
+        args.aug_rotate,
+        args.aug_jitter_coords,
+        args.aug_shift_coords,
+        args.aug_flip,
+        args.aug_mask_feats,
+        args.aug_jitter_length,
+        args.aug_elasticate,
+    ]
+    aug_fns = [
+        None,
+        RandomScaleCoords(p=0.2),
+        RandomRotate(p=0.5),
+        RandomJitter(p=0.2),
+        RandomShift(p=0.2),
+        RandomFlip(p=1),
+        RandomMaskFeats(p=0.2),
+        RandomJitterLength(p=0.2),
+        RandomElasticate(p=0.2),
+    ]
+    feat_augs = [aug_fns[i] for i in range(len(aug_switchs)) if aug_switchs[i] and aug_fns[i] is not None]
+    feat_augs = Compose(feat_augs) if feat_augs else None
+
+    topo_aug_switchs = [
+        args.aug_drop_tree,
+        args.aug_skip_parent_node,
+        args.aug_swap_sibling_subtrees,
+    ]
+    topo_aug_fns = [
+        RandomDropSubTrees(probs=[0.05], max_cnt=5),
+        RandomSkipParentNode(probs=[0.05], max_cnt=10),
+        RandomSwapSiblingSubTrees(probs=[0.05], max_cnt=10),
+    ]
+    topo_augs = [topo_aug_fns[i] for i in range(len(topo_aug_switchs)) if topo_aug_switchs[i]]
+    topo_augs = Compose(topo_augs) if topo_augs else None
+
+    # Persistence augmentation
+    persistence_aug = None
+    if args.use_persistence_aug:
+        persistence_aug = CombinedPersistenceAugmentation(
+            translation_scale=args.pers_translation_scale,
+            noise_scale=args.pers_noise_scale,
+            sigma_min=args.pers_sigma_min,
+            sigma_max=args.pers_sigma_max,
+        )
+
+    # Create datasets
+    collate_fn = get_collate_fn(device, use_images=True)
+
+    trainset = NeuronTreeDataset(
+        phase="train",
+        dataset=args.dataset,
+        label_dict=LABEL_DICT[args.dataset],
+        input_features=args.input_features,
+        topology_transformations=topo_augs,
+        attribute_transformations=feat_augs,
+        use_images=True,
+        image_size=args.image_size,
+        cache_images=args.cache_images,
+        persistence_augmentation=persistence_aug,
+        sigma_px=args.sigma_px,
+    )
+
+    testset = NeuronTreeDataset(
+        phase="test",
+        dataset=args.dataset,
+        label_dict=LABEL_DICT[args.dataset],
+        input_features=args.input_features,
+        use_images=True,
+        image_size=args.image_size,
+        cache_images=args.cache_images,
+    )
+
+    train_loader = DataLoader(
+        dataset=trainset,
+        batch_size=args.batch_size,
+        collate_fn=collate_fn,
+        shuffle=True,
+        num_workers=4,
+        pin_memory=True,
+    )
+
+    test_loader = DataLoader(
+        dataset=testset,
+        batch_size=args.batch_size,
+        collate_fn=collate_fn,
+        shuffle=False,
+        num_workers=4,
+        pin_memory=True,
+    )
+
+    logger.info(f"Train samples: {len(trainset)}, Test samples: {len(testset)}")
+    logger.info(f"Number of classes: {len(trainset.classes)}")
+
+    # Create fine-tuning model
+    logger.info("=> Creating fine-tuning model...")
+    model = FineTuneModel(
+        pretrained_model=pretrained_model,
+        num_classes=len(trainset.classes),
+        mode=args.mode,
+        freeze_encoders=args.freeze_encoders,
+        fusion_mode=args.fusion_mode if args.mode == "multimodal" else None,
+        dropout=args.dropout,
+        label_smoothing=args.label_smoothing,
+        use_projection=args.use_projection,
+        use_arcface=args.use_arcface,
+        arcface_s=args.arcface_s,
+        arcface_m=args.arcface_m,
+        freeze_image_only=args.freeze_image_only
+    ).to(device)
+
+    del pretrained_model
+    logger.info(model)
+
+    # Count parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    logger.info(f"Total parameters: {total_params:,}")
+    logger.info(f"Trainable parameters: {trainable_params:,}")
+
+    # Optimizer
+    optimizer = torch.optim.AdamW(
+        filter(lambda p: p.requires_grad, model.parameters()),
+        lr=args.lr,
+        weight_decay=args.wd,
+    )
+
+    cudnn.benchmark = True
+
+    # Scheduler
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+        optimizer,
+        T_0=25,
+        T_mult=1,
+        eta_min=args.lr * 0.01
+    )
+
+    # Training loop
+    best_metric = 0.0
+    best_epoch = 0
+    patience_counter = 0
+    total_iters = len(train_loader) * args.epochs
+    current_iter = 0
+    start_time = time.time()
+
+    logger.info("=> Starting fine-tuning...")
+    logger.info(f"=> Regularization: dropout={args.dropout}, label_smoothing={args.label_smoothing}")
+
+    for epoch in range(args.start_epoch + 1, args.epochs + 1):
+
+        if epoch == args.linear_probe_epochs + 1 and args.linear_probe_epochs > 0:
+            logger.info("="*30)
+            logger.info(f"Linear probe phase completed. Unfreezing encoders for full fine-tuning.")
+            model.unfreeze_encoders()
+            
+            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+            logger.info(f"Trainable parameters after unfreezing: {trainable_params:,}")
+
+            optimizer = torch.optim.AdamW(
+                filter(lambda p: p.requires_grad, model.parameters()),
+                lr=args.lr / 10,
+                weight_decay=args.wd,
+            )
+            logger.info(f"Optimizer reset to include unfrozen parameters with a new learning rate of {args.lr / 10}.")
+            logger.info("="*30)
+
+        model.train()
+        epoch_loss = 0.0
+        correct = 0
+        total = 0
+
+        for step, batch in enumerate(train_loader):
+            try:
+                batch = batch.to(device)
+
+                if args.mixup_alpha > 0:
+                    _, _, features = model(batch, return_features=True)
+                    labels = batch.label.cuda() if not batch.label.is_cuda else batch.label
+
+                    features, labels_a, labels_b, lam = mixup_data(features, labels, args.mixup_alpha)
+
+                    if model.use_arcface:
+                        extracted_features = model.feature_extractor(features)
+                        logits = model.arcface(extracted_features, labels)
+                        loss = model.criterion(logits, labels)
+                    else:
+                        logits = model.classifier(features)
+                        loss = mixup_criterion(model.criterion, logits, labels_a, labels_b, lam)
+                else:
+                    loss, logits = model(batch)
+
+                optimizer.zero_grad()
+                loss.backward()
+
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
+                optimizer.step()
+
+                pred = logits.argmax(dim=1)
+                labels = batch.label.cuda() if not batch.label.is_cuda else batch.label
+                correct += (pred == labels).sum().item()
+                total += labels.size(0)
+
+                epoch_loss += loss.item()
+                current_iter += 1
+
+                if step % 10 == 0:
+                    current_time = time.time()
+                    elapsed = current_time - start_time
+
+                    log_str = (
+                        f"Epoch {epoch:03d} | Step {step:03d}/{len(train_loader)} | "
+                        f"Loss {loss.item():.4f} | "
+                        f"LR {optimizer.param_groups[0]['lr']:.6f} | "
+                        f"Elapsed {str(datetime.timedelta(seconds=int(elapsed)))}"
+                    )
+                    logger.info(log_str)
+            except Exception as e:
+                logger.info(f"Error in step {step}: {e}")
+                continue
+
+        scheduler.step()
+
+        avg_loss = epoch_loss / len(train_loader)
+        train_acc = correct / total * 100
+        logger.info(f"Epoch {epoch:03d} | Avg Loss: {avg_loss:.4f} | Train Acc: {train_acc:.2f}%")
+
+        # Evaluation
+        if epoch % args.val_freq == 0:
+            logger.info("=> Evaluating...")
+
+            if args.eval_mode == "accuracy":
+                test_acc = evaluate_accuracy(model, test_loader, device)
+                logger.info(f"  Test Accuracy: {test_acc:.2f}%")
+                metric = test_acc
+            else:
+                knn_acc = evaluate_knn(model, train_loader, test_loader, device, args.knn_k)
+                logger.info(f"  KNN Accuracy (k={args.knn_k}): {knn_acc:.2f}%")
+                metric = knn_acc
+
+            if metric > best_metric:
+                best_metric = metric
+                best_epoch = epoch
+                patience_counter = 0
+
+                checkpoint_path = f"{args.work_dir}/best_model.pth"
+                save_checkpoint(
+                    {
+                        "epoch": epoch,
+                        "state_dict": model.state_dict(),
+                        "optimizer": optimizer.state_dict(),
+                        "metric": metric,
+                        "mode": args.mode,
+                    },
+                    is_best=True,
+                    filename=checkpoint_path,
+                )
+                logger.info(f"  Saved new best checkpoint: {checkpoint_path}")
+            else:
+                patience_counter += 1
+
+            logger.info(f"  Best: {best_metric:.2f}% at epoch {best_epoch}")
+
+            if args.early_stopping_patience > 0 and patience_counter >= args.early_stopping_patience:
+                logger.info(f"  Early stopping triggered")
+                break
+
+        # Save periodic checkpoint
+        if epoch % args.save_freq == 0:
+            checkpoint_path = f"{args.work_dir}/epoch_{epoch}.pth"
+            save_checkpoint(
+                {
+                    "epoch": epoch,
+                    "state_dict": model.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                },
+                is_best=False,
+                filename=checkpoint_path,
+            )
+            logger.info(f"Saved checkpoint: {checkpoint_path}")
+
+    logger.info("Fine-tuning complete!")
+    logger.info(f"Best {args.eval_mode}: {best_metric:.2f}% at epoch {best_epoch}")

graphformer/__init__.py

@@ -0,0 +1,12 @@
+"""
+GraPHFormer: Graph-Persistence Hybrid Transformer for Neuron Morphology
+
+A CLIP-style contrastive learning framework for neuron representation learning
+that combines tree-structured morphology with persistence images.
+"""
+
+from .models import CLIPModel, FineTuneModel, CLIPLoss
+from .losses import InfoNCELoss, NTXentLoss, TripletLoss
+
+__version__ = "1.0.0"
+__all__ = ["CLIPModel", "FineTuneModel", "CLIPLoss", "InfoNCELoss", "NTXentLoss", "TripletLoss"]

graphformer/augmentations/__init__.py

@@ -0,0 +1,23 @@
+"""Augmentation strategies for GraPHFormer."""
+
+from .tree_augmentations import (
+    Compose,
+    RandomDropSubTrees, RandomSkipParentNode, RandomSwapSiblingSubTrees,
+    RandomRotate, RandomJitter, RandomShift, RandomFlip,
+    RandomScaleCoords, RandomScaleFeats, RandomMaskFeats,
+    RandomElasticate, RandomJitterLength,
+)
+from .persistence_augmentations import (
+    PersistenceSpaceAugmentation, SigmaVariationAugmentation,
+    CombinedPersistenceAugmentation, get_default_augmentation,
+)
+
+__all__ = [
+    "Compose",
+    "RandomDropSubTrees", "RandomSkipParentNode", "RandomSwapSiblingSubTrees",
+    "RandomRotate", "RandomJitter", "RandomShift", "RandomFlip",
+    "RandomScaleCoords", "RandomScaleFeats", "RandomMaskFeats",
+    "RandomElasticate", "RandomJitterLength",
+    "PersistenceSpaceAugmentation", "SigmaVariationAugmentation",
+    "CombinedPersistenceAugmentation", "get_default_augmentation",
+]

graphformer/augmentations/persistence_augmentations.py

@@ -0,0 +1,304 @@
+"""
+Persistence Image Augmentation Module
+
+This module provides augmentation strategies specifically designed for persistence images
+generated from neuron morphology data. Unlike standard image augmentations (ColorJitter, etc.),
+these augmentations work in the persistence space (birth, persistence) to create meaningful
+variations while preserving the underlying topological structure.
+
+Augmentation Strategies:
+1. Translation in birth/persistence space - shifts points in the diagram
+2. Gaussian noise addition - adds small random perturbations to point positions
+3. Sigma variation - varies the Gaussian kernel width during image generation
+"""
+
+import numpy as np
+import torch
+from typing import Tuple, Optional
+import random
+
+
+class PersistenceSpaceAugmentation:
+    """
+    Augmentation that operates in persistence space before image generation.
+
+    This augmentation modifies the birth/persistence coordinates of points
+    in the persistence diagram, then regenerates the image with these modified coordinates.
+    """
+
+    def __init__(
+        self,
+        translation_scale: float = 0.05,
+        noise_scale: float = 0.02,
+        translation_prob: float = 0.5,
+        noise_prob: float = 0.5,
+        persistence_scale_prob: float = 0.0,
+        persistence_scale_min: float = 0.9,
+        persistence_scale_max: float = 1.1,
+        radius_perturb_prob: float = 0.0,
+        radius_perturb_min: float = 0.85,
+        radius_perturb_max: float = 1.15,
+    ):
+        """
+        Args:
+            translation_scale: Scale of random translation (relative to data range)
+            noise_scale: Scale of Gaussian noise (relative to data range)
+            translation_prob: Probability of applying translation
+            noise_prob: Probability of applying Gaussian noise
+            persistence_scale_prob: Probability of applying persistence scaling
+            persistence_scale_min: Minimum scaling factor for persistence
+            persistence_scale_max: Maximum scaling factor for persistence
+            radius_perturb_prob: Probability of applying radius perturbation
+            radius_perturb_min: Minimum scaling factor for radius
+            radius_perturb_max: Maximum scaling factor for radius
+        """
+        self.translation_scale = translation_scale
+        self.noise_scale = noise_scale
+        self.translation_prob = translation_prob
+        self.noise_prob = noise_prob
+        self.persistence_scale_prob = persistence_scale_prob
+        self.persistence_scale_min = persistence_scale_min
+        self.persistence_scale_max = persistence_scale_max
+        self.radius_perturb_prob = radius_perturb_prob
+        self.radius_perturb_min = radius_perturb_min
+        self.radius_perturb_max = radius_perturb_max
+
+    def augment_pairs_features(self, pairs_feats, global_bounds=None):
+        """
+        Augment the pairs features in persistence space.
+
+        Args:
+            pairs_feats: List of dictionaries with 'birth', 'death', 'persistence', 'mean_radius'
+            global_bounds: Tuple of (birth_min, birth_max, pers_min, pers_max)
+
+        Returns:
+            Augmented pairs_feats
+        """
+        if not pairs_feats or len(pairs_feats) == 0:
+            return pairs_feats
+
+        # Create a copy to avoid modifying original
+        augmented_feats = [f.copy() for f in pairs_feats]
+
+        # Extract births and persistence values
+        births = np.array([f['birth'] for f in augmented_feats])
+        pers = np.array([f['persistence'] for f in augmented_feats])
+
+        # Determine data range for scaling
+        if global_bounds is not None:
+            birth_range = global_bounds[1] - global_bounds[0]
+            pers_range = global_bounds[3] - global_bounds[2]
+        else:
+            birth_range = births.max() - births.min() + 1e-6
+            pers_range = pers.max() - pers.min() + 1e-6
+
+        # Apply random translation in birth/persistence space
+        if random.random() < self.translation_prob:
+            birth_shift = np.random.uniform(-self.translation_scale, self.translation_scale) * birth_range
+            pers_shift = np.random.uniform(-self.translation_scale, self.translation_scale) * pers_range
+
+            births += birth_shift
+            pers += pers_shift
+
+        # Apply Gaussian noise
+        if random.random() < self.noise_prob:
+            birth_noise = np.random.normal(0, self.noise_scale * birth_range, size=births.shape)
+            pers_noise = np.random.normal(0, self.noise_scale * pers_range, size=pers.shape)
+
+            births += birth_noise
+            pers += pers_noise
+
+        # Apply persistence scaling
+        if random.random() < self.persistence_scale_prob:
+            alpha = np.random.uniform(self.persistence_scale_min, self.persistence_scale_max)
+            pers *= alpha
+
+        # Apply radius perturbation
+        if random.random() < self.radius_perturb_prob:
+            beta = np.random.uniform(self.radius_perturb_min, self.radius_perturb_max)
+            for f in augmented_feats:
+                if 'mean_radius' in f:
+                    f['mean_radius'] *= beta
+
+        # Ensure persistence values remain positive
+        pers = np.maximum(pers, 1e-9)
+
+        # Update the augmented features
+        for i, f in enumerate(augmented_feats):
+            f['birth'] = float(births[i])
+            f['persistence'] = float(pers[i])
+            # Note: death = birth - persistence (in TMD convention where persistence is negative)
+            # Actually in this code: persistence = birth - death, so death = birth - persistence
+            f['death'] = float(births[i] - pers[i])
+
+        return augmented_feats
+
+
+class SigmaVariationAugmentation:
+    """
+    Augmentation that varies the sigma parameter during Gaussian kernel application.
+
+    This creates different "blur" levels in the persistence image, which can help
+    the model learn features at multiple scales.
+    """
+
+    def __init__(
+        self,
+        sigma_min: float = 12.0,
+        sigma_max: float = 20.0,
+        prob: float = 1.0,
+    ):
+        """
+        Args:
+            sigma_min: Minimum sigma value
+            sigma_max: Maximum sigma value
+            prob: Probability of varying sigma (1.0 means always vary)
+        """
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.prob = prob
+
+    def sample_sigma(self, base_sigma: float = 16.0) -> float:
+        """
+        Sample a sigma value.
+
+        Args:
+            base_sigma: Base sigma value (unused when prob=1.0, kept for compatibility)
+
+        Returns:
+            Sampled sigma value
+        """
+        if random.random() < self.prob:
+            return np.random.uniform(self.sigma_min, self.sigma_max)
+        else:
+            return base_sigma
+
+
+class CombinedPersistenceAugmentation:
+    """
+    Combines multiple persistence space augmentations.
+
+    This is the main augmentation class that should be used for training.
+    """
+
+    def __init__(
+        self,
+        translation_scale: float = 0.05,
+        noise_scale: float = 0.02,
+        sigma_min: float = 12.0,
+        sigma_max: float = 20.0,
+        translation_prob: float = 0.5,
+        noise_prob: float = 0.5,
+        sigma_variation_prob: float = 1.0,
+        persistence_scale_prob: float = 0.0,
+        persistence_scale_min: float = 0.9,
+        persistence_scale_max: float = 1.1,
+        radius_perturb_prob: float = 0.0,
+        radius_perturb_min: float = 0.85,
+        radius_perturb_max: float = 1.15,
+    ):
+        """
+        Args:
+            translation_scale: Scale of random translation in birth/persistence space
+            noise_scale: Scale of Gaussian noise
+            sigma_min: Minimum sigma for Gaussian kernel
+            sigma_max: Maximum sigma for Gaussian kernel
+            translation_prob: Probability of applying translation
+            noise_prob: Probability of applying noise
+            sigma_variation_prob: Probability of varying sigma
+            persistence_scale_prob: Probability of applying persistence scaling
+            persistence_scale_min: Minimum scaling factor for persistence
+            persistence_scale_max: Maximum scaling factor for persistence
+            radius_perturb_prob: Probability of applying radius perturbation
+            radius_perturb_min: Minimum scaling factor for radius
+            radius_perturb_max: Maximum scaling factor for radius
+        """
+        self.space_aug = PersistenceSpaceAugmentation(
+            translation_scale=translation_scale,
+            noise_scale=noise_scale,
+            translation_prob=translation_prob,
+            noise_prob=noise_prob,
+            persistence_scale_prob=persistence_scale_prob,
+            persistence_scale_min=persistence_scale_min,
+            persistence_scale_max=persistence_scale_max,
+            radius_perturb_prob=radius_perturb_prob,
+            radius_perturb_min=radius_perturb_min,
+            radius_perturb_max=radius_perturb_max,
+        )
+        self.sigma_aug = SigmaVariationAugmentation(
+            sigma_min=sigma_min,
+            sigma_max=sigma_max,
+            prob=sigma_variation_prob,
+        )
+
+    def augment_pairs_features(self, pairs_feats, global_bounds=None):
+        """Augment pairs features in persistence space."""
+        return self.space_aug.augment_pairs_features(pairs_feats, global_bounds)
+
+    def sample_sigma(self, base_sigma: float = 16.0) -> float:
+        """Sample a sigma value for image generation."""
+        return self.sigma_aug.sample_sigma(base_sigma)
+
+
+def get_default_augmentation(mode='train'):
+    """
+    Get default augmentation configuration.
+
+    Args:
+        mode: 'train' or 'test'
+
+    Returns:
+        CombinedPersistenceAugmentation instance
+    """
+    if mode == 'train':
+        return CombinedPersistenceAugmentation(
+            translation_scale=0.05,  # 5% of range
+            noise_scale=0.02,         # 2% of range
+            sigma_min=12.0,
+            sigma_max=20.0,
+            translation_prob=0.5,
+            noise_prob=0.5,
+            sigma_variation_prob=1.0,
+        )
+    else:
+        # No augmentation for test
+        return CombinedPersistenceAugmentation(
+            translation_scale=0.0,
+            noise_scale=0.0,
+            sigma_min=16.0,
+            sigma_max=16.0,
+            translation_prob=0.0,
+            noise_prob=0.0,
+            sigma_variation_prob=0.0,
+        )
+
+
+# Example usage
+if __name__ == "__main__":
+    # Create augmentation
+    aug = get_default_augmentation('train')
+
+    # Example pairs features
+    pairs_feats = [
+        {'birth': 100.0, 'death': 50.0, 'persistence': 50.0, 'mean_radius': 2.5},
+        {'birth': 150.0, 'death': 80.0, 'persistence': 70.0, 'mean_radius': 3.0},
+        {'birth': 200.0, 'death': 120.0, 'persistence': 80.0, 'mean_radius': 2.8},
+    ]
+
+    # Augment
+    global_bounds = (0, 300, 0, 100)
+    augmented = aug.augment_pairs_features(pairs_feats, global_bounds)
+
+    print("Original pairs:")
+    for f in pairs_feats[:2]:
+        print(f"  birth={f['birth']:.2f}, persistence={f['persistence']:.2f}")
+
+    print("\nAugmented pairs:")
+    for f in augmented[:2]:
+        print(f"  birth={f['birth']:.2f}, persistence={f['persistence']:.2f}")
+
+    # Sample sigma values
+    print("\nSampled sigma values:")
+    for _ in range(5):
+        sigma = aug.sample_sigma()
+        print(f"  sigma={sigma:.2f}")

graphformer/augmentations/tree_augmentations.py

@@ -0,0 +1,435 @@
+import numpy as np
+from nltk.tree import Tree
+import random
+
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, coords_and_feats):
+        for t in self.transforms:
+            coords_and_feats = t(coords_and_feats)
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        sep = "\n\t"
+        msg = "Compose of ["
+        for t in self.transforms:
+            msg += f"{sep}{t},"
+        msg += "]\n"
+        return msg
+
+    def __repr__(self) -> str:
+        sep = "\n\t"
+        msg = "Compose of ["
+        for t in self.transforms:
+            msg += f"{sep}{t}"
+        msg += "]\n"
+        return msg
+
+
+class RandomDropSubTrees(object):
+    def __init__(self, probs=[0.1, 0.2, 0.3, 0.4, 0.5], max_cnt=5):
+        self.probs = probs
+        self.max_cnt = max_cnt
+        self.cnt = 0
+
+    def remove_subtrees(self, root, level_idx):
+        reduced_root = Tree(root.label(), [])
+        if len(root) == 0:
+            return reduced_root
+        p = np.random.uniform(0, 1, len(root))
+        # padding the dropping probabilities
+        if level_idx >= len(self.probs):
+            level_idx = len(self.probs) - 1
+        for child_idx in range(len(root)):
+            if self.cnt > self.max_cnt:
+                reduced_root.append(root[child_idx])
+                continue
+            else:
+                if p[child_idx] > self.probs[level_idx]:
+                    reduced_root.append(
+                        self.remove_subtrees(root[child_idx], level_idx + 1)
+                    )
+                else:
+                    self.cnt += 1
+        return reduced_root
+
+    def __call__(self, tree):
+        self.cnt = 0
+        return self.remove_subtrees(tree, level_idx=0)
+
+    def __str__(self) -> str:
+        return f"RandomDropSubTrees(probs={self.probs}, max_cnt={self.max_cnt})"
+
+    def __repr__(self) -> str:
+        return f"RandomDropSubTrees(probs={self.probs}, max_cnt={self.max_cnt})"
+
+
+class RandomSkipParentNode(object):
+    def __init__(self, probs=[0.1], max_cnt=5):
+        self.probs = probs
+        self.max_cnt = max_cnt
+        self.cnt = 0
+
+    def move_grandson_to_son(self, root, level_idx):
+        if len(root) == 0 or len(root) == 1:
+            return root
+        p = np.random.uniform(0, 1, len(root))
+        # padding the dropping probabilities
+        if level_idx >= len(self.probs):
+            level_idx = len(self.probs) - 1
+        for child_idx in range(len(root)):
+            if self.cnt >= self.max_cnt:
+                break
+            if p[child_idx] < self.probs[level_idx]:
+                if len(root[child_idx]) == 0 or len(root[child_idx]) == 1:
+                    continue
+                else:
+                    idx = random.randint(0, len(root[child_idx]) - 1)
+                    root[child_idx] = root[child_idx][idx]
+                    self.cnt += 1
+            else:
+                root[child_idx] = self.move_grandson_to_son(
+                    root[child_idx], level_idx + 1
+                )
+        return root
+
+    def __call__(self, tree):
+        self.cnt = 0
+        return self.move_grandson_to_son(tree, level_idx=0)
+
+    def __str__(self) -> str:
+        return f"RandomSkipParentNode(probs={self.probs}, max_cnt={self.max_cnt})"
+
+    def __repr__(self) -> str:
+        return f"RandomSkipParentNode(probs={self.probs}, max_cnt={self.max_cnt})"
+
+
+class RandomSwapSiblingSubTrees(object):
+    def __init__(self, probs=[0.1], max_cnt=5):
+        self.probs = probs
+        self.max_cnt = max_cnt
+        self.cnt = 0
+
+    def swap_sibling_subtrees(self, root, level_idx):
+        if len(root) < 2:
+            return root
+        p = np.random.uniform(0, 1, len(root))
+        # padding the dropping probabilities
+        if level_idx >= len(self.probs):
+            level_idx = len(self.probs) - 1
+        for child_idx in range(len(root)):
+            if self.cnt >= self.max_cnt:
+                break
+            if p[child_idx] < self.probs[level_idx]:
+                if len(root[child_idx]) < 2:
+                    continue
+                else:
+                    my_subtree_idx = random.randint(0, len(root[child_idx]) - 1)
+                    sibling_idx = random.randint(0, len(root) - 1)
+                    if len(root[sibling_idx]) == 0:
+                        continue
+                    sibling_subtree_idx = random.randint(0, len(root[sibling_idx]) - 1)
+                    my_subtree = root[child_idx][my_subtree_idx].copy()
+                    sibling_subtree = root[sibling_idx][sibling_subtree_idx].copy()
+                    root[child_idx][my_subtree_idx] = sibling_subtree
+                    root[sibling_idx][sibling_subtree_idx] = my_subtree
+                    self.cnt += 1
+            else:
+                root[child_idx] = self.swap_sibling_subtrees(
+                    root[child_idx], level_idx + 1
+                )
+        return root
+
+    def __call__(self, tree):
+        self.cnt = 0
+        return self.swap_sibling_subtrees(tree, level_idx=0)
+
+    def __str__(self) -> str:
+        return f"RandomSwapSiblingSubTrees(probs={self.probs}, max_cnt={self.max_cnt})"
+
+    def __repr__(self) -> str:
+        return f"RandomSwapSiblingSubTrees(probs={self.probs}, max_cnt={self.max_cnt})"
+
+
+class RandomRotateAligned(object):
+    def __init__(self, p=0.5, axis=2):
+        self.prob = p
+        self.axis = axis
+
+    def __call__(self, coords_and_feats):
+        coord = coords_and_feats[:, :3]
+        if np.random.rand() < self.prob:
+            angle = np.random.uniform() * 2 * np.pi
+            cos, sin = np.cos(angle), np.sin(angle)
+            R_x = np.array([[1, 0, 0], [0, cos, -sin], [0, sin, cos]])
+            R_y = np.array([[cos, 0, sin], [0, 1, 0], [-sin, 0, cos]])
+            R_z = np.array([[cos, -sin, 0], [sin, cos, 0], [0, 0, 1]])
+            R = [R_x, R_y, R_z][self.axis]
+            coord = np.dot(coord, R)
+        coords_and_feats[:, :3] = coord
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomRotateAligned(p={self.prob},axis={self.axis})"
+
+    def __repr__(self) -> str:
+        return f"RandomRotateAligned(p={self.prob},axis={self.axis})"
+
+
+class RandomRotate(object):
+    def __init__(self, sigma=0.03, clip=0.09, p=0.5):
+        self.sigma = sigma
+        self.clip = clip
+        self.prob = p
+
+    def __call__(self, coords_and_feats):
+        coord = coords_and_feats[:, :3]
+        if np.random.rand() < self.prob:
+            angle_x = np.random.uniform() * 2 * np.pi
+            angle_y = np.random.uniform() * 2 * np.pi
+            angle_z = np.random.uniform() * 2 * np.pi
+            cos_x, sin_x = np.cos(angle_x), np.sin(angle_x)
+            cos_y, sin_y = np.cos(angle_y), np.sin(angle_y)
+            cos_z, sin_z = np.cos(angle_z), np.sin(angle_z)
+            R_x = np.array([[1, 0, 0], [0, cos_x, -sin_x], [0, sin_x, cos_x]])
+            R_y = np.array([[cos_y, 0, sin_y], [0, 1, 0], [-sin_y, 0, cos_y]])
+            R_z = np.array([[cos_z, -sin_z, 0], [sin_z, cos_z, 0], [0, 0, 1]])
+            R = np.dot(R_z, np.dot(R_y, R_x))
+            coord = np.dot(coord, R)
+        coords_and_feats[:, :3] = coord
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomRotate(p={self.prob})"
+
+    def __repr__(self) -> str:
+        return f"RandomRotate(p={self.prob})"
+
+
+class RandomMaskFeats(object):
+    def __init__(self, p=0.2):
+        self.prob = p
+
+    def __call__(self, coords_and_feats):
+        if len(coords_and_feats[0]) > 5:
+            feats = coords_and_feats[:, 5:]
+            feats[
+                :,
+                np.random.choice(
+                    np.arange(len(feats[0])), int(len(feats[0]) * self.prob)
+                ),
+            ] = 0
+            coords_and_feats[:, 5:] = feats
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomMaskFeats(p={self.prob})"
+
+    def __repr__(self) -> str:
+        return f"RandomMaskFeats(p={self.prob})"
+
+
+class RandomElasticate(object):
+    def __init__(self, p=0.2, scales=[0.8, 1.2]):
+        self.prob = p
+        self.scales = scales
+
+    def __call__(self, coords_and_feats):
+        if len(coords_and_feats[0]) > 5:
+            if np.random.rand() < self.prob:
+                branches = coords_and_feats[:, 5:]
+                scales = np.random.uniform(
+                    self.scales[0], self.scales[1], branches.shape
+                )
+                branches *= scales
+                coords_and_feats[:, 5:] = branches
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomElasticate(p={self.prob}, scales={self.scales})"
+
+    def __repr__(self) -> str:
+        return f"RandomElasticate(p={self.prob}, scales={self.scales})"
+
+
+class RandomScaleCoords(object):
+    def __init__(self, scale=[0.8, 1.2], p=0.5):
+        self.scale = scale
+        self.prob = p
+
+    def __call__(self, coords_and_feats):
+        if np.random.rand() < self.prob:
+            scale = np.random.uniform(self.scale[0], self.scale[1])
+            coords_and_feats[:, :4] *= scale
+            if len(coords_and_feats[0]) > 5:
+                coords_and_feats[:, 5:] *= scale
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomScaleCoords(p={self.prob}, scale={self.scale})"
+
+    def __repr__(self) -> str:
+        return f"RandomScaleCoords(p={self.prob}, scale={self.scale})"
+
+
+class RandomScaleCoordsTranslation(object):
+    def __init__(self, scale=[0.5, 2], p=0.5):
+        self.scale = scale
+        self.prob = p
+
+    def __call__(self, coords_and_feats):
+        if np.random.rand() < self.prob:
+            scale = np.random.uniform(self.scale[0], self.scale[1])
+            coord1 = coords_and_feats[:, :4]
+            coord1 *= scale
+            coords_and_feats[:, :4] = coord1
+            if len(coords_and_feats[0]) > 5:
+                coord2 = coords_and_feats[:, 5:]
+                coord2 *= scale
+                coords_and_feats[:, 5:] = coord2
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomScaleCoordsTranslation(p={self.prob}, scale={self.scale})"
+
+    def __repr__(self) -> str:
+        return f"RandomScaleCoordsTranslation(p={self.prob}, scale={self.scale})"
+
+
+class RandomScaleFeats(object):
+    def __init__(self, scale=[0.5, 2], p=0.5):
+        self.scale = scale
+        self.prob = p
+
+    def __call__(self, coords_and_feats):
+        feats = coords_and_feats[:, 4:]
+        if np.random.rand() < self.prob:
+            scale = np.random.uniform(self.scale[0], self.scale[1])
+            feats *= scale
+        coords_and_feats[:, 4:] = feats
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomScaleFeats(p={self.prob}, scale={self.scale})"
+
+    def __repr__(self) -> str:
+        return f"RandomScaleFeats(p={self.prob}, scale={self.scale})"
+
+
+class RandomShift(object):
+    def __init__(self, shift=[5, 5, 5], p=0.5):
+        self.shift = shift
+        self.prob = p
+
+    def __call__(self, coords_and_feats):
+        coord = coords_and_feats[:, :3]
+        if np.random.rand() < self.prob:
+            shift_x = np.random.uniform(-self.shift[0], self.shift[0])
+            shift_y = np.random.uniform(-self.shift[1], self.shift[1])
+            shift_z = np.random.uniform(-self.shift[2], self.shift[2])
+            coord += [shift_x, shift_y, shift_z]
+        coords_and_feats[:, :3] = coord
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomShift(p={self.prob}, shift={self.shift})"
+
+    def __repr__(self) -> str:
+        return f"RandomShift(p={self.prob}, shift={self.shift})"
+
+
+class RandomFlip(object):
+    def __init__(self, p=0.5):
+        self.prob = p
+
+    def __call__(self, coords_and_feats):
+        coord = coords_and_feats[:, :3]
+        if np.random.rand() < self.prob:
+            if np.random.rand() < 0.5:
+                coord[:, 0] = -coord[:, 0]
+            if np.random.rand() < 0.5:
+                coord[:, 1] = -coord[:, 1]
+        coords_and_feats[:, :3] = coord
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomFlip(p={self.prob})"
+
+    def __repr__(self) -> str:
+        return f"RandomFlip(p={self.prob})"
+
+
+class RandomJitter(object):
+    def __init__(self, sigma=1, clip=5, p=0.5):
+        self.sigma = sigma
+        self.clip = clip
+        self.prob = p
+
+    def __call__(self, coords_and_feats):
+        coord = coords_and_feats[:, :3]
+        assert self.clip > 0
+        if np.random.rand() < self.prob:
+            jitter = np.clip(
+                self.sigma * np.random.randn(coord.shape[0], 3), -self.clip, self.clip
+            )
+            coord += jitter
+        coords_and_feats[:, :3] = coord
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return f"RandomJitter(p={self.prob}, sigma={self.sigma}, clip={self.clip})"
+
+    def __repr__(self) -> str:
+        return f"RandomJitter(p={self.prob}, sigma={self.sigma}, clip={self.clip})"
+
+
+class RandomJitterLength(object):
+    def __init__(self, sigma=0.1, clip=1, p=0.5):
+        self.sigma = sigma
+        self.clip = clip
+        self.prob = p
+
+    def __call__(self, coords_and_feats):
+        feats1 = coords_and_feats[:, 3:4]
+        assert self.clip > 0
+        if np.random.rand() < self.prob:
+            jitter1 = np.clip(
+                self.sigma * np.random.randn(*feats1.shape), -self.clip, self.clip
+            )
+            feats1 += jitter1
+            if len(coords_and_feats[0]) > 5:
+                feats2 = coords_and_feats[:, 5:]
+                jitter2 = np.clip(
+                    self.sigma * np.random.randn(*feats2.shape), -self.clip, self.clip
+                )
+                feats2 += jitter2
+                coords_and_feats[:, 5:] = feats2
+        coords_and_feats[:, 3:4] = feats1
+        return coords_and_feats
+
+    def __str__(self) -> str:
+        return (
+            f"RandomJitterLength(p={self.prob}, sigma={self.sigma}, clip={self.clip})"
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"RandomJitterLength(p={self.prob}, sigma={self.sigma}, clip={self.clip})"
+        )
+
+
+if __name__ == "__main__":
+    transform = RandomScaleCoordsTranslation(p=1)
+
+    coords_feats = np.random.rand(1024, 3)
+    # coords_feats = np.concatenate([
+    #     np.zeros((1,29)),
+    #     coords_feats
+    # ])
+    print(coords_feats[-1])
+    transformed = transform(coords_feats)
+    print(transformed[-1])
+    print(coords_feats[-1])

graphformer/losses/__init__.py

@@ -0,0 +1,9 @@
+"""Loss functions for GraPHFormer."""
+
+from .infonce import InfoNCELoss, SymmetricInfoNCELoss, HardNegativeInfoNCELoss, MultiModalInfoNCELoss
+from .contrastive import NTXentLoss, TripletLoss, CombinedContrastiveLoss
+
+__all__ = [
+    "InfoNCELoss", "SymmetricInfoNCELoss", "HardNegativeInfoNCELoss", "MultiModalInfoNCELoss",
+    "NTXentLoss", "TripletLoss", "CombinedContrastiveLoss",
+]

graphformer/losses/contrastive.py

@@ -0,0 +1,442 @@
+"""
+Contrastive Loss Implementations for Neuron Representation Learning
+
+Includes:
+- NT-Xent (Normalized Temperature-scaled Cross Entropy) - SimCLR loss
+- Triplet Loss (with various mining strategies)
+- InfoNCE variants (imported from infonce_loss.py)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class NTXentLoss(nn.Module):
+    """
+    NT-Xent (Normalized Temperature-scaled Cross Entropy) Loss
+
+    Used in SimCLR (Chen et al., 2020) - "A Simple Framework for Contrastive Learning"
+
+    Key features:
+    - Normalizes embeddings to unit sphere
+    - Uses temperature scaling
+    - Treats all other samples in batch as negatives
+    - Symmetric loss over both views
+    """
+    def __init__(self, temperature=0.5, use_cosine_similarity=True):
+        """
+        Args:
+            temperature: Temperature parameter for scaling (default 0.5 from SimCLR)
+            use_cosine_similarity: If True, use cosine similarity (default True)
+        """
+        super(NTXentLoss, self).__init__()
+        self.temperature = temperature
+        self.use_cosine_similarity = use_cosine_similarity
+
+    def forward(self, z_i, z_j):
+        """
+        Compute NT-Xent loss between two views
+
+        Args:
+            z_i: (B, D) embeddings from view 1
+            z_j: (B, D) embeddings from view 2
+
+        Returns:
+            loss: scalar NT-Xent loss
+        """
+        batch_size = z_i.shape[0]
+
+        # Normalize embeddings
+        z_i = F.normalize(z_i, dim=1)
+        z_j = F.normalize(z_j, dim=1)
+
+        # Concatenate both views: (2B, D)
+        z = torch.cat([z_i, z_j], dim=0)
+
+        # Compute similarity matrix: (2B, 2B)
+        if self.use_cosine_similarity:
+            similarity_matrix = torch.matmul(z, z.T)
+        else:
+            # Dot product similarity
+            similarity_matrix = torch.matmul(z, z.T)
+
+        # Remove diagonal elements (self-similarity)
+        # Create mask to exclude diagonal
+        mask = torch.eye(2 * batch_size, dtype=torch.bool, device=z.device)
+        similarity_matrix = similarity_matrix.masked_fill(mask, -9e15)
+
+        # Apply temperature scaling
+        similarity_matrix = similarity_matrix / self.temperature
+
+        # Positive pairs are at positions (i, i+B) and (i+B, i)
+        # For each sample, its positive is at distance B
+        positive_samples = torch.cat([
+            torch.arange(batch_size, 2 * batch_size),  # Positives for first half
+            torch.arange(0, batch_size)                 # Positives for second half
+        ], dim=0).to(z.device)
+
+        # Extract positive similarities
+        positive_sim = similarity_matrix[torch.arange(2 * batch_size), positive_samples]
+        positive_sim = positive_sim.reshape(2 * batch_size, 1)
+
+        # Negative similarities: all other samples (already in similarity_matrix)
+        # For numerical stability, use log-sum-exp trick
+        negatives_mask = ~mask
+        negatives_mask[torch.arange(2 * batch_size), positive_samples] = False
+
+        # Compute loss using LogSumExp
+        # loss = -log(exp(pos) / (exp(pos) + sum(exp(neg))))
+        # For numerical stability: -pos + log(exp(pos) + sum(exp(neg)))
+
+        # Get all similarities (positive at correct position + negatives)
+        all_similarities = similarity_matrix  # (2B, 2B)
+
+        # Create labels for cross-entropy: positive is at position batch_size or 0
+        labels = positive_samples
+
+        # Use cross-entropy loss (equivalent to NT-Xent)
+        loss = F.cross_entropy(all_similarities, labels)
+
+        return loss
+
+
+class TripletLoss(nn.Module):
+    """
+    Triplet Loss for metric learning
+
+    L = max(d(a, p) - d(a, n) + margin, 0)
+
+    Where:
+    - a: anchor
+    - p: positive (same class as anchor)
+    - n: negative (different class)
+    - d: distance metric (L2 or cosine)
+    """
+    def __init__(self, margin=1.0, distance_metric='euclidean', mining='batch_hard'):
+        """
+        Args:
+            margin: Margin for triplet loss
+            distance_metric: 'euclidean' or 'cosine'
+            mining: 'batch_hard', 'batch_all', or 'semi_hard'
+        """
+        super(TripletLoss, self).__init__()
+        self.margin = margin
+        self.distance_metric = distance_metric
+        self.mining = mining
+
+    def _compute_distance_matrix(self, embeddings):
+        """Compute pairwise distance matrix."""
+        if self.distance_metric == 'euclidean':
+            # Efficient euclidean distance: ||a-b||^2 = ||a||^2 + ||b||^2 - 2*a·b
+            dot_product = torch.matmul(embeddings, embeddings.T)
+            square_norm = torch.diag(dot_product)
+            distances = square_norm.unsqueeze(1) - 2.0 * dot_product + square_norm.unsqueeze(0)
+            distances = torch.clamp(distances, min=0.0)  # Numerical stability
+            distances = torch.sqrt(distances + 1e-16)
+        elif self.distance_metric == 'cosine':
+            # Cosine distance = 1 - cosine_similarity
+            embeddings_norm = F.normalize(embeddings, dim=1)
+            cosine_sim = torch.matmul(embeddings_norm, embeddings_norm.T)
+            distances = 1.0 - cosine_sim
+        else:
+            raise ValueError(f"Unknown distance metric: {self.distance_metric}")
+
+        return distances
+
+    def _batch_hard_mining(self, embeddings, labels):
+        """
+        Batch hard mining: for each anchor, select hardest positive and hardest negative.
+
+        Hardest positive: furthest positive sample
+        Hardest negative: closest negative sample
+        """
+        batch_size = embeddings.shape[0]
+
+        # Compute pairwise distances
+        distances = self._compute_distance_matrix(embeddings)
+
+        # Create masks for positives and negatives
+        labels = labels.unsqueeze(1)
+        positive_mask = (labels == labels.T).float()
+        negative_mask = (labels != labels.T).float()
+
+        # Remove self-comparisons from positive mask
+        positive_mask = positive_mask - torch.eye(batch_size, device=embeddings.device)
+
+        # Hard positive: maximum distance among positives
+        # Set non-positive distances to 0 so they won't be selected
+        positive_distances = distances * positive_mask
+        hardest_positive_dist, _ = torch.max(positive_distances, dim=1)
+
+        # Hard negative: minimum distance among negatives
+        # Set non-negative distances to large value so they won't be selected
+        negative_distances = distances + (1.0 - negative_mask) * 1e9
+        hardest_negative_dist, _ = torch.min(negative_distances, dim=1)
+
+        # Compute triplet loss
+        triplet_loss = F.relu(hardest_positive_dist - hardest_negative_dist + self.margin)
+
+        return triplet_loss.mean()
+
+    def _batch_all_mining(self, embeddings, labels):
+        """
+        Batch all mining: use all valid triplets in the batch.
+        """
+        batch_size = embeddings.shape[0]
+
+        # Compute pairwise distances
+        distances = self._compute_distance_matrix(embeddings)
+
+        # Create masks
+        labels = labels.unsqueeze(1)
+        positive_mask = (labels == labels.T).float()
+        negative_mask = (labels != labels.T).float()
+
+        # Remove self-comparisons
+        positive_mask = positive_mask - torch.eye(batch_size, device=embeddings.device)
+
+        # Get anchor-positive distances: (B, B)
+        anchor_positive_dist = distances.unsqueeze(2)  # (B, B, 1)
+
+        # Get anchor-negative distances: (B, B)
+        anchor_negative_dist = distances.unsqueeze(1)  # (B, 1, B)
+
+        # Compute triplet loss for all valid triplets
+        triplet_loss = anchor_positive_dist - anchor_negative_dist + self.margin
+
+        # Mask out invalid triplets
+        # Valid triplet: (i, j, k) where label[i] == label[j] != label[k]
+        valid_triplets = positive_mask.unsqueeze(2) * negative_mask.unsqueeze(1)
+
+        # Apply mask and ReLU
+        triplet_loss = triplet_loss * valid_triplets
+        triplet_loss = F.relu(triplet_loss)
+
+        # Count valid triplets
+        num_valid = valid_triplets.sum()
+
+        if num_valid > 0:
+            triplet_loss = triplet_loss.sum() / num_valid
+        else:
+            triplet_loss = torch.tensor(0.0, device=embeddings.device)
+
+        return triplet_loss
+
+    def _semi_hard_mining(self, embeddings, labels):
+        """
+        Semi-hard mining: select negatives that are harder than positive but still within margin.
+
+        Semi-hard negative: d(a,p) < d(a,n) < d(a,p) + margin
+        """
+        batch_size = embeddings.shape[0]
+
+        # Compute pairwise distances
+        distances = self._compute_distance_matrix(embeddings)
+
+        # Create masks
+        labels = labels.unsqueeze(1)
+        positive_mask = (labels == labels.T).float()
+        negative_mask = (labels != labels.T).float()
+
+        # Remove self-comparisons
+        positive_mask = positive_mask - torch.eye(batch_size, device=embeddings.device)
+
+        losses = []
+
+        for i in range(batch_size):
+            # Get positive distances for anchor i
+            pos_dists = distances[i] * positive_mask[i]
+            if pos_dists.sum() == 0:
+                continue
+
+            # Select a positive (use hardest for stability)
+            pos_dist = pos_dists.max()
+
+            # Get negative distances for anchor i
+            neg_dists = distances[i] * negative_mask[i]
+
+            # Semi-hard negatives: pos_dist < neg_dist < pos_dist + margin
+            semi_hard_mask = (neg_dists > pos_dist) & (neg_dists < pos_dist + self.margin)
+
+            if semi_hard_mask.any():
+                # Use hardest semi-hard negative (closest to anchor)
+                semi_hard_negatives = neg_dists.clone()
+                semi_hard_negatives[~semi_hard_mask] = 1e9
+                neg_dist = semi_hard_negatives.min()
+            else:
+                # Fall back to hardest negative
+                neg_dists_masked = neg_dists + (1.0 - negative_mask[i]) * 1e9
+                neg_dist = neg_dists_masked.min()
+
+            # Compute triplet loss
+            loss = F.relu(pos_dist - neg_dist + self.margin)
+            losses.append(loss)
+
+        if len(losses) > 0:
+            return torch.stack(losses).mean()
+        else:
+            return torch.tensor(0.0, device=embeddings.device)
+
+    def forward(self, embeddings, labels):
+        """
+        Compute triplet loss
+
+        Args:
+            embeddings: (B, D) embeddings
+            labels: (B,) class labels for each sample
+
+        Returns:
+            loss: scalar triplet loss
+        """
+        if self.mining == 'batch_hard':
+            return self._batch_hard_mining(embeddings, labels)
+        elif self.mining == 'batch_all':
+            return self._batch_all_mining(embeddings, labels)
+        elif self.mining == 'semi_hard':
+            return self._semi_hard_mining(embeddings, labels)
+        else:
+            raise ValueError(f"Unknown mining strategy: {self.mining}")
+
+
+class CombinedContrastiveLoss(nn.Module):
+    """
+    Combine multiple contrastive losses with configurable weights
+
+    Example: NT-Xent + Triplet Loss
+    """
+    def __init__(self, loss_types=['ntxent'], loss_weights=None, **loss_kwargs):
+        """
+        Args:
+            loss_types: List of loss types ('ntxent', 'triplet', 'infonce')
+            loss_weights: List of weights for each loss (default: equal weights)
+            **loss_kwargs: Keyword arguments for each loss
+                - ntxent_temperature: Temperature for NT-Xent
+                - triplet_margin: Margin for Triplet Loss
+                - triplet_mining: Mining strategy for Triplet Loss
+        """
+        super(CombinedContrastiveLoss, self).__init__()
+        self.loss_types = loss_types
+
+        if loss_weights is None:
+            self.loss_weights = [1.0] * len(loss_types)
+        else:
+            self.loss_weights = loss_weights
+
+        # Initialize losses
+        self.losses = nn.ModuleDict()
+
+        for loss_type in loss_types:
+            if loss_type == 'ntxent':
+                temp = loss_kwargs.get('ntxent_temperature', 0.5)
+                self.losses['ntxent'] = NTXentLoss(temperature=temp)
+            elif loss_type == 'triplet':
+                margin = loss_kwargs.get('triplet_margin', 1.0)
+                mining = loss_kwargs.get('triplet_mining', 'batch_hard')
+                distance = loss_kwargs.get('triplet_distance', 'euclidean')
+                self.losses['triplet'] = TripletLoss(
+                    margin=margin,
+                    distance_metric=distance,
+                    mining=mining
+                )
+            else:
+                raise ValueError(f"Unknown loss type: {loss_type}")
+
+    def forward(self, z_i, z_j, labels=None):
+        """
+        Compute combined loss
+
+        Args:
+            z_i: (B, D) embeddings from view 1
+            z_j: (B, D) embeddings from view 2
+            labels: (B,) class labels (required for triplet loss)
+
+        Returns:
+            loss: scalar combined loss
+            loss_dict: dictionary with individual loss values
+        """
+        total_loss = 0.0
+        loss_dict = {}
+
+        for i, loss_type in enumerate(self.loss_types):
+            weight = self.loss_weights[i]
+
+            if loss_type == 'ntxent':
+                loss_val = self.losses['ntxent'](z_i, z_j)
+            elif loss_type == 'triplet':
+                if labels is None:
+                    raise ValueError("Triplet loss requires labels")
+                # Combine both views for triplet loss
+                embeddings = torch.cat([z_i, z_j], dim=0)
+                combined_labels = torch.cat([labels, labels], dim=0)
+                loss_val = self.losses['triplet'](embeddings, combined_labels)
+
+            loss_dict[loss_type] = loss_val.item()
+            total_loss += weight * loss_val
+
+        loss_dict['total'] = total_loss.item()
+
+        return total_loss, loss_dict
+
+
+if __name__ == "__main__":
+    print("Testing Contrastive Loss Implementations...")
+    print("=" * 60)
+
+    B, D = 32, 128
+    num_classes = 10
+
+    # Test 1: NT-Xent Loss
+    print("\n1. NT-Xent Loss (SimCLR):")
+    ntxent = NTXentLoss(temperature=0.5)
+
+    z_i = torch.randn(B, D)
+    z_j = torch.randn(B, D)
+
+    loss = ntxent(z_i, z_j)
+    print(f"   Loss: {loss.item():.4f}")
+
+    # Test 2: Triplet Loss (Batch Hard)
+    print("\n2. Triplet Loss (Batch Hard Mining):")
+    triplet = TripletLoss(margin=1.0, mining='batch_hard')
+
+    embeddings = torch.randn(B, D)
+    labels = torch.randint(0, num_classes, (B,))
+
+    loss = triplet(embeddings, labels)
+    print(f"   Loss: {loss.item():.4f}")
+
+    # Test 3: Triplet Loss (Batch All)
+    print("\n3. Triplet Loss (Batch All Mining):")
+    triplet_all = TripletLoss(margin=1.0, mining='batch_all')
+
+    loss = triplet_all(embeddings, labels)
+    print(f"   Loss: {loss.item():.4f}")
+
+    # Test 4: Triplet Loss (Semi-Hard)
+    print("\n4. Triplet Loss (Semi-Hard Mining):")
+    triplet_semi = TripletLoss(margin=1.0, mining='semi_hard')
+
+    loss = triplet_semi(embeddings, labels)
+    print(f"   Loss: {loss.item():.4f}")
+
+    # Test 5: Combined Loss (NT-Xent + Triplet)
+    print("\n5. Combined Loss (NT-Xent + Triplet):")
+    combined = CombinedContrastiveLoss(
+        loss_types=['ntxent', 'triplet'],
+        loss_weights=[1.0, 0.5],
+        ntxent_temperature=0.5,
+        triplet_margin=1.0,
+        triplet_mining='batch_hard'
+    )
+
+    z_i = torch.randn(B, D)
+    z_j = torch.randn(B, D)
+    labels = torch.randint(0, num_classes, (B,))
+
+    loss, loss_dict = combined(z_i, z_j, labels)
+    print(f"   Total Loss: {loss.item():.4f}")
+    print(f"   Loss breakdown: {loss_dict}")
+
+    print("\n" + "=" * 60)
+    print("All tests passed!")

graphformer/losses/infonce.py

@@ -0,0 +1,303 @@
+"""
+InfoNCE Loss Implementation
+
+InfoNCE (Information Noise-Contrastive Estimation) loss from:
+"Representation Learning with Contrastive Predictive Coding" (van den Oord et al., 2018)
+
+Used in many contrastive learning methods like MoCo, SimCLR, etc.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class InfoNCELoss(nn.Module):
+    """
+    InfoNCE loss for contrastive learning
+
+    The loss encourages positive pairs to have high similarity while
+    negative pairs have low similarity using a contrastive objective.
+    """
+    def __init__(self, temperature=0.07, reduction='mean'):
+        """
+        Args:
+            temperature: Temperature parameter for scaling similarities
+            reduction: 'mean' or 'sum' for loss reduction
+        """
+        super(InfoNCELoss, self).__init__()
+        self.temperature = temperature
+        self.reduction = reduction
+
+    def forward(self, anchor, positive, negatives=None):
+        """
+        Compute InfoNCE loss
+
+        Args:
+            anchor: (B, D) anchor embeddings
+            positive: (B, D) positive embeddings (paired with anchor)
+            negatives: (B, N, D) or None - negative embeddings
+                      If None, uses all other samples in batch as negatives (in-batch negatives)
+
+        Returns:
+            loss: scalar InfoNCE loss
+        """
+        # Normalize embeddings
+        anchor = F.normalize(anchor, dim=-1)
+        positive = F.normalize(positive, dim=-1)
+
+        B = anchor.shape[0]
+
+        if negatives is None:
+            # Use in-batch negatives (like SimCLR)
+            # Each sample uses all other samples as negatives
+
+            # Positive similarity: (B,)
+            pos_sim = torch.sum(anchor * positive, dim=-1) / self.temperature
+
+            # Negative similarities: (B, B)
+            # anchor vs all positives (including itself)
+            neg_sim = torch.matmul(anchor, positive.T) / self.temperature
+
+            # Create labels (diagonal elements are positive pairs)
+            labels = torch.arange(B, device=anchor.device)
+
+            # InfoNCE loss using cross-entropy
+            loss = F.cross_entropy(neg_sim, labels, reduction=self.reduction)
+
+        else:
+            # Explicit negatives provided
+            negatives = F.normalize(negatives, dim=-1)
+
+            # Positive similarity: (B,)
+            pos_sim = torch.sum(anchor * positive, dim=-1, keepdim=True) / self.temperature
+
+            # Negative similarities: (B, N)
+            neg_sim = torch.matmul(anchor.unsqueeze(1), negatives.transpose(1, 2)).squeeze(1) / self.temperature
+
+            # Concatenate positive and negative similarities
+            logits = torch.cat([pos_sim, neg_sim], dim=1)  # (B, 1+N)
+
+            # Labels: first position (index 0) is the positive
+            labels = torch.zeros(B, dtype=torch.long, device=anchor.device)
+
+            # InfoNCE loss
+            loss = F.cross_entropy(logits, labels, reduction=self.reduction)
+
+        return loss
+
+
+class SymmetricInfoNCELoss(nn.Module):
+    """
+    Symmetric InfoNCE loss (like CLIP)
+
+    Computes InfoNCE in both directions:
+    - anchor -> positive (image -> text in CLIP)
+    - positive -> anchor (text -> image in CLIP)
+
+    Then averages the two losses.
+    """
+    def __init__(self, temperature=0.07):
+        super(SymmetricInfoNCELoss, self).__init__()
+        self.temperature = temperature
+
+    def forward(self, embedding_a, embedding_b):
+        """
+        Compute symmetric InfoNCE loss between two sets of embeddings
+
+        Args:
+            embedding_a: (B, D) - first modality (e.g., tree embeddings)
+            embedding_b: (B, D) - second modality (e.g., image embeddings)
+
+        Returns:
+            loss: scalar symmetric InfoNCE loss
+        """
+        # Normalize
+        embedding_a = F.normalize(embedding_a, dim=-1)
+        embedding_b = F.normalize(embedding_b, dim=-1)
+
+        B = embedding_a.shape[0]
+
+        # Compute similarity matrix: (B, B)
+        similarity = torch.matmul(embedding_a, embedding_b.T) / self.temperature
+
+        # Labels: diagonal elements are positive pairs
+        labels = torch.arange(B, device=embedding_a.device)
+
+        # Loss in both directions
+        loss_a2b = F.cross_entropy(similarity, labels)
+        loss_b2a = F.cross_entropy(similarity.T, labels)
+
+        # Average
+        loss = (loss_a2b + loss_b2a) / 2
+
+        return loss
+
+
+class HardNegativeInfoNCELoss(nn.Module):
+    """
+    InfoNCE with hard negative mining
+
+    Selects the hardest negatives (highest similarity) for each anchor
+    to make training more challenging and effective.
+    """
+    def __init__(self, temperature=0.07, num_hard_negatives=10):
+        """
+        Args:
+            temperature: Temperature parameter
+            num_hard_negatives: Number of hard negatives to mine per anchor
+        """
+        super(HardNegativeInfoNCELoss, self).__init__()
+        self.temperature = temperature
+        self.num_hard_negatives = num_hard_negatives
+
+    def forward(self, anchor, positive, negative_pool):
+        """
+        Compute InfoNCE with hard negative mining
+
+        Args:
+            anchor: (B, D) anchor embeddings
+            positive: (B, D) positive embeddings
+            negative_pool: (M, D) pool of negative embeddings (M >> B)
+
+        Returns:
+            loss: scalar InfoNCE loss with hard negatives
+        """
+        # Normalize
+        anchor = F.normalize(anchor, dim=-1)
+        positive = F.normalize(positive, dim=-1)
+        negative_pool = F.normalize(negative_pool, dim=-1)
+
+        B = anchor.shape[0]
+
+        # Positive similarity: (B,)
+        pos_sim = torch.sum(anchor * positive, dim=-1, keepdim=True) / self.temperature
+
+        # Compute similarity to all negatives: (B, M)
+        all_neg_sim = torch.matmul(anchor, negative_pool.T) / self.temperature
+
+        # Select top-k hard negatives (highest similarity = hardest)
+        hard_neg_sim, _ = torch.topk(all_neg_sim, k=self.num_hard_negatives, dim=1)
+
+        # Concatenate positive and hard negative similarities
+        logits = torch.cat([pos_sim, hard_neg_sim], dim=1)  # (B, 1+K)
+
+        # Labels: first position is positive
+        labels = torch.zeros(B, dtype=torch.long, device=anchor.device)
+
+        # InfoNCE loss
+        loss = F.cross_entropy(logits, labels)
+
+        return loss
+
+
+class MultiModalInfoNCELoss(nn.Module):
+    """
+    Multi-modal InfoNCE for learning joint embeddings across multiple modalities
+
+    Example: tree structure + persistence image + graph features
+    """
+    def __init__(self, temperature=0.07, weight_modalities=None):
+        """
+        Args:
+            temperature: Temperature parameter
+            weight_modalities: List of weights for each modality pair loss
+        """
+        super(MultiModalInfoNCELoss, self).__init__()
+        self.temperature = temperature
+        self.weight_modalities = weight_modalities
+
+    def forward(self, embeddings):
+        """
+        Compute InfoNCE across all modality pairs
+
+        Args:
+            embeddings: List of (B, D) embeddings for each modality
+
+        Returns:
+            loss: scalar multi-modal InfoNCE loss
+        """
+        num_modalities = len(embeddings)
+
+        # Normalize all embeddings
+        embeddings = [F.normalize(emb, dim=-1) for emb in embeddings]
+
+        B = embeddings[0].shape[0]
+        labels = torch.arange(B, device=embeddings[0].device)
+
+        # Compute loss for all pairs of modalities
+        total_loss = 0
+        num_pairs = 0
+
+        for i in range(num_modalities):
+            for j in range(i + 1, num_modalities):
+                # Similarity matrix
+                similarity = torch.matmul(embeddings[i], embeddings[j].T) / self.temperature
+
+                # Symmetric loss
+                loss_ij = F.cross_entropy(similarity, labels)
+                loss_ji = F.cross_entropy(similarity.T, labels)
+                pair_loss = (loss_ij + loss_ji) / 2
+
+                # Apply weight if provided
+                if self.weight_modalities is not None:
+                    weight = self.weight_modalities[num_pairs]
+                    pair_loss = weight * pair_loss
+
+                total_loss += pair_loss
+                num_pairs += 1
+
+        # Average over all pairs
+        loss = total_loss / num_pairs
+
+        return loss
+
+
+if __name__ == "__main__":
+    print("Testing InfoNCE Loss implementations...")
+
+    B, D = 32, 128
+
+    # Test 1: Basic InfoNCE with in-batch negatives
+    print("\n1. Basic InfoNCE Loss:")
+    loss_fn = InfoNCELoss(temperature=0.07)
+
+    anchor = torch.randn(B, D)
+    positive = torch.randn(B, D)
+
+    loss = loss_fn(anchor, positive)
+    print(f"   Loss: {loss.item():.4f}")
+
+    # Test 2: Symmetric InfoNCE (CLIP-style)
+    print("\n2. Symmetric InfoNCE Loss (CLIP-style):")
+    symmetric_loss = SymmetricInfoNCELoss(temperature=0.07)
+
+    tree_embeddings = torch.randn(B, D)
+    image_embeddings = torch.randn(B, D)
+
+    loss = symmetric_loss(tree_embeddings, image_embeddings)
+    print(f"   Loss: {loss.item():.4f}")
+
+    # Test 3: Hard Negative Mining
+    print("\n3. InfoNCE with Hard Negative Mining:")
+    hard_neg_loss = HardNegativeInfoNCELoss(temperature=0.07, num_hard_negatives=10)
+
+    anchor = torch.randn(B, D)
+    positive = torch.randn(B, D)
+    negative_pool = torch.randn(500, D)  # Large pool of negatives
+
+    loss = hard_neg_loss(anchor, positive, negative_pool)
+    print(f"   Loss: {loss.item():.4f}")
+
+    # Test 4: Multi-modal InfoNCE
+    print("\n4. Multi-modal InfoNCE:")
+    multimodal_loss = MultiModalInfoNCELoss(temperature=0.07)
+
+    tree_emb = torch.randn(B, D)
+    image_emb = torch.randn(B, D)
+    graph_emb = torch.randn(B, D)
+
+    loss = multimodal_loss([tree_emb, image_emb, graph_emb])
+    print(f"   Loss: {loss.item():.4f}")
+
+    print("\nAll tests passed!")

graphformer/models/__init__.py

@@ -0,0 +1,15 @@
+"""Model components for GraPHFormer."""
+
+from .clip_model import CLIPModel, CLIPLoss
+from .finetune_model import FineTuneModel, ArcMarginProduct
+from .tree_encoder import TreeLSTM, TreeLSTMv2, TreeLSTMDouble, TreeLSTMCell
+from .image_encoder import ImageEncoder, SimpleCNN, SmallViT, PersistenceViT, DINOv2ImageEncoder
+from .fusion import CrossAttentionFusion, BiDirectionalCrossAttention, GatedFusion, CMF, MultiHeadCrossModalAttention
+
+__all__ = [
+    "CLIPModel", "CLIPLoss",
+    "FineTuneModel", "ArcMarginProduct",
+    "TreeLSTM", "TreeLSTMv2", "TreeLSTMDouble", "TreeLSTMCell",
+    "ImageEncoder", "SimpleCNN", "SmallViT", "PersistenceViT", "DINOv2ImageEncoder",
+    "CrossAttentionFusion", "BiDirectionalCrossAttention", "GatedFusion", "CMF", "MultiHeadCrossModalAttention",
+]

graphformer/models/clip_model.py

@@ -0,0 +1,170 @@
+"""
+CLIP-style Contrastive Model for Neuron Morphology
+
+Aligns tree structure representations with persistence images using
+contrastive learning with separate encoders.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .tree_encoder import TreeLSTM, TreeLSTMv2, TreeLSTMDouble
+from .image_encoder import ImageEncoder
+
+
+class CLIPLoss(nn.Module):
+    """CLIP-style symmetric contrastive loss"""
+    def __init__(self, temperature=0.07):
+        super(CLIPLoss, self).__init__()
+        self.temperature = temperature
+
+    def forward(self, tree_features, image_features):
+        """
+        Args:
+            tree_features: (B, dim) - normalized tree embeddings
+            image_features: (B, dim) - normalized image embeddings
+        Returns:
+            loss: scalar contrastive loss
+        """
+        tree_features = F.normalize(tree_features, dim=-1)
+        image_features = F.normalize(image_features, dim=-1)
+
+        logits = torch.matmul(tree_features, image_features.T) / self.temperature
+
+        batch_size = tree_features.shape[0]
+        labels = torch.arange(batch_size, device=tree_features.device)
+
+        loss_tree_to_image = F.cross_entropy(logits, labels)
+        loss_image_to_tree = F.cross_entropy(logits.T, labels)
+
+        loss = (loss_tree_to_image + loss_image_to_tree) / 2
+
+        return loss
+
+
+class CLIPModel(nn.Module):
+    """CLIP-style model with separate tree and image encoders"""
+    def __init__(self, args):
+        super(CLIPModel, self).__init__()
+
+        self.tree_encoder_type = args.tree_model
+
+        # Tree encoder
+        if args.tree_model == "ori":
+            self.tree_encoder = TreeLSTM(
+                x_size=len(args.input_features),
+                h_size=args.h_size,
+                num_classes=0,
+                fc=False,
+                bn=args.bn,
+                mode=args.child_mode,
+            )
+        elif args.tree_model == "v2":
+            self.tree_encoder = TreeLSTMv2(
+                x_size=len(args.input_features),
+                h_size=args.h_size,
+                num_classes=0,
+                fc=False,
+                bn=args.bn,
+                mode=args.child_mode,
+            )
+        elif args.tree_model == "double":
+            self.tree_encoder = TreeLSTMDouble(
+                x_size=len(args.input_features),
+                h_size=args.h_size,
+                num_classes=0,
+                fc=False,
+                bn=args.bn,
+                mode=args.child_mode,
+            )
+        else:
+            raise ValueError(f"Unknown tree model: {args.tree_model}")
+
+        # Image encoder
+        self.image_encoder = ImageEncoder(
+            output_dim=args.h_size,
+            model_type=args.image_encoder,
+            image_size=args.image_size,
+            freeze_backbone=args.freeze_image_backbone,
+        )
+
+        # Get image encoder output dimension
+        if hasattr(self.image_encoder.encoder if hasattr(self.image_encoder, 'encoder') else self.image_encoder, 'feat_dim'):
+            image_feat_dim = self.image_encoder.encoder.feat_dim if hasattr(self.image_encoder, 'encoder') else self.image_encoder.feat_dim
+        else:
+            image_feat_dim = args.h_size
+
+        # Projection heads
+        if getattr(args, 'single_linear_proj', False):
+            self.tree_projection = nn.Linear(args.h_size, args.embed_dim)
+            self.image_projection = nn.Linear(image_feat_dim, args.embed_dim)
+        else:
+            self.tree_projection = nn.Sequential(
+                nn.Linear(args.h_size, args.h_size),
+                nn.ReLU(),
+                nn.Linear(args.h_size, args.embed_dim),
+            )
+            self.image_projection = nn.Sequential(
+                nn.Linear(image_feat_dim, args.h_size),
+                nn.ReLU(),
+                nn.Linear(args.h_size, args.embed_dim),
+            )
+
+        self.loss_type = args.loss_type
+
+        # Loss function
+        if args.loss_type == 'infonce':
+            from ..losses import SymmetricInfoNCELoss
+            self.criterion = SymmetricInfoNCELoss(temperature=args.temperature)
+        elif args.loss_type == 'ntxent':
+            from ..losses import NTXentLoss
+            self.criterion = NTXentLoss(temperature=args.temperature)
+        elif args.loss_type == 'triplet':
+            from ..losses import TripletLoss
+            self.criterion = TripletLoss(
+                margin=args.triplet_margin,
+                distance_metric=args.triplet_distance,
+                mining=args.triplet_mining
+            )
+        else:  # 'clip' (default)
+            self.criterion = CLIPLoss(temperature=args.temperature)
+
+    def encode_tree(self, batch):
+        """Encode tree data"""
+        tree_feats = self.tree_encoder(batch)
+        tree_embed = self.tree_projection(tree_feats)
+        return tree_embed
+
+    def encode_image(self, images):
+        """Encode image data"""
+        image_feats = self.image_encoder(images)
+        image_embed = self.image_projection(image_feats)
+        return image_embed
+
+    def forward(self, batch, return_recon=False):
+        """
+        Forward pass computing CLIP loss
+
+        Args:
+            batch: contains batch.graph, batch.feats, batch.images
+            return_recon: unused, kept for compatibility
+        Returns:
+            loss: Contrastive loss
+        """
+        tree_embed = self.encode_tree(batch)
+        images = batch.images.cuda() if not batch.images.is_cuda else batch.images
+
+        image_feats = self.image_encoder(images)
+        image_embed = self.image_projection(image_feats)
+
+        # Contrastive loss
+        if self.loss_type == 'triplet':
+            embeddings = torch.cat([tree_embed, image_embed], dim=0)
+            labels_gpu = batch.label.cuda() if not batch.label.is_cuda else batch.label
+            labels = torch.cat([labels_gpu, labels_gpu], dim=0)
+            clip_loss = self.criterion(embeddings, labels)
+        else:
+            clip_loss = self.criterion(tree_embed, image_embed)
+
+        return clip_loss

graphformer/models/finetune_model.py

@@ -0,0 +1,322 @@
+"""
+Fine-tuning Model for GraPHFormer
+
+Supports three modes: image_only, tree_only, multimodal
+Loads pretrained weights from CLIP-style training checkpoints
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .fusion import (
+    CrossAttentionFusion, BiDirectionalCrossAttention,
+    GatedFusion, CMF, MultiHeadCrossModalAttention
+)
+
+
+class ArcMarginProduct(nn.Module):
+    """ArcFace: Additive Angular Margin Loss"""
+    def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False):
+        super(ArcMarginProduct, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.s = s
+        self.m = m
+        self.easy_margin = easy_margin
+
+        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
+        nn.init.xavier_uniform_(self.weight)
+
+        self.cos_m = np.cos(m)
+        self.sin_m = np.sin(m)
+        self.th = np.cos(np.pi - m)
+        self.mm = np.sin(np.pi - m) * m
+
+    def forward(self, input, label):
+        cosine = F.linear(F.normalize(input), F.normalize(self.weight))
+        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
+
+        phi = cosine * self.cos_m - sine * self.sin_m
+
+        if self.easy_margin:
+            phi = torch.where(cosine > 0, phi, cosine)
+        else:
+            phi = torch.where(cosine > self.th, phi, cosine - self.mm)
+
+        one_hot = torch.zeros(cosine.size(), device=input.device)
+        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
+
+        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
+        output *= self.s
+
+        return output
+
+
+class FineTuneModel(nn.Module):
+    """Fine-tuning model with classification head"""
+    def __init__(
+        self,
+        pretrained_model,
+        num_classes,
+        mode='multimodal',
+        freeze_encoders=False,
+        fusion_mode='concat',
+        dropout=0.5,
+        label_smoothing=0.0,
+        use_projection=False,
+        use_arcface=False,
+        arcface_s=30.0,
+        arcface_m=0.50,
+        freeze_image_only=False
+    ):
+        """
+        Args:
+            pretrained_model: Pretrained CLIPModel
+            num_classes: Number of classes for classification
+            mode: 'image_only', 'tree_only', or 'multimodal'
+            freeze_encoders: If True, freeze encoder weights
+            fusion_mode: For multimodal - 'concat', 'add', 'cross_attention', 'bi_attention', 'gated', 'cmf', 'mhcma'
+            dropout: Dropout rate for first layer
+            label_smoothing: Label smoothing factor
+            use_projection: If True, use projection heads from pretrained model
+            use_arcface: If True, use ArcFace loss instead of CrossEntropy
+            arcface_s: ArcFace scale parameter
+            arcface_m: ArcFace margin parameter
+            freeze_image_only: If True, freeze only image encoder
+        """
+        super(FineTuneModel, self).__init__()
+
+        self.mode = mode
+        self.fusion_mode = fusion_mode
+        self.use_projection = use_projection
+        self.use_arcface = use_arcface
+        self.tree_encoder_type = pretrained_model.tree_encoder_type
+
+        # Copy encoders from pretrained model
+        if mode in ['tree_only', 'multimodal']:
+            self.tree_encoder = pretrained_model.tree_encoder
+            if use_projection:
+                self.tree_projection = pretrained_model.tree_projection
+
+        if mode in ['image_only', 'multimodal']:
+            self.image_encoder = pretrained_model.image_encoder
+            if use_projection:
+                self.image_projection = pretrained_model.image_projection
+
+        # Freeze encoders if requested
+        if freeze_encoders:
+            if mode in ['tree_only', 'multimodal']:
+                for param in self.tree_encoder.parameters():
+                    param.requires_grad = False
+                if use_projection:
+                    for param in self.tree_projection.parameters():
+                        param.requires_grad = False
+
+            if mode in ['image_only', 'multimodal']:
+                for param in self.image_encoder.parameters():
+                    param.requires_grad = False
+                if use_projection:
+                    for param in self.image_projection.parameters():
+                        param.requires_grad = False
+
+        # Freeze only image encoder
+        if freeze_image_only and mode == 'multimodal':
+            for param in self.image_encoder.parameters():
+                param.requires_grad = False
+            if use_projection:
+                for param in self.image_projection.parameters():
+                    param.requires_grad = False
+
+        # Get embedding dimensions
+        if use_projection:
+            tree_embed_dim = pretrained_model.tree_projection[-1].out_features
+            image_embed_dim = pretrained_model.image_projection[-1].out_features
+        else:
+            if mode in ['tree_only', 'multimodal']:
+                tree_embed_dim = pretrained_model.tree_encoder.h_size
+            else:
+                tree_embed_dim = 0
+
+            if mode in ['image_only', 'multimodal']:
+                if hasattr(pretrained_model.image_encoder, 'encoder'):
+                    if hasattr(pretrained_model.image_encoder.encoder, 'feat_dim'):
+                        image_embed_dim = pretrained_model.image_encoder.encoder.feat_dim
+                    else:
+                        image_embed_dim = pretrained_model.image_encoder.encoder[-1].in_features
+                elif hasattr(pretrained_model.image_encoder, 'feat_dim'):
+                    image_embed_dim = pretrained_model.image_encoder.feat_dim
+                else:
+                    image_embed_dim = pretrained_model.tree_encoder.h_size
+            else:
+                image_embed_dim = 0
+
+        # Setup fusion for multimodal
+        if mode == 'multimodal':
+            if fusion_mode == 'concat':
+                fusion_dim = tree_embed_dim + image_embed_dim
+            elif fusion_mode == 'add':
+                fusion_dim = min(tree_embed_dim, image_embed_dim)
+                if tree_embed_dim != image_embed_dim:
+                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
+                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
+            elif fusion_mode == 'cross_attention':
+                fusion_dim = min(tree_embed_dim, image_embed_dim)
+                if tree_embed_dim != image_embed_dim:
+                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
+                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
+                self.fusion_layer = CrossAttentionFusion(fusion_dim, num_heads=4)
+            elif fusion_mode == 'bi_attention':
+                fusion_dim = min(tree_embed_dim, image_embed_dim)
+                if tree_embed_dim != image_embed_dim:
+                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
+                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
+                self.fusion_layer = BiDirectionalCrossAttention(fusion_dim, num_heads=4)
+            elif fusion_mode == 'gated':
+                fusion_dim = min(tree_embed_dim, image_embed_dim)
+                if tree_embed_dim != image_embed_dim:
+                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
+                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
+                self.fusion_layer = GatedFusion(fusion_dim)
+            elif fusion_mode == 'cmf':
+                fusion_dim = min(tree_embed_dim, image_embed_dim)
+                if tree_embed_dim != image_embed_dim:
+                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
+                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
+                self.fusion_layer = CMF(fusion_dim)
+            elif fusion_mode == 'mhcma':
+                fusion_dim = min(tree_embed_dim, image_embed_dim)
+                if tree_embed_dim != image_embed_dim:
+                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
+                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
+                self.fusion_layer = MultiHeadCrossModalAttention(fusion_dim, num_heads=8)
+            else:
+                raise ValueError(f"Unknown fusion_mode: {fusion_mode}")
+            input_dim = fusion_dim
+        else:
+            if mode == 'tree_only':
+                input_dim = tree_embed_dim
+            else:
+                input_dim = image_embed_dim
+
+        # Classification head
+        if use_arcface:
+            self.feature_extractor = nn.Sequential(
+                nn.Dropout(dropout),
+                nn.Linear(input_dim, input_dim // 2),
+                nn.BatchNorm1d(input_dim // 2),
+                nn.ReLU(),
+                nn.Dropout(dropout * 0.7),
+            )
+            self.arcface = ArcMarginProduct(input_dim // 2, num_classes, s=arcface_s, m=arcface_m)
+            self.classifier = None
+        else:
+            self.classifier = nn.Sequential(
+                nn.Dropout(dropout),
+                nn.Linear(input_dim, input_dim // 2),
+                nn.BatchNorm1d(input_dim // 2),
+                nn.ReLU(),
+                nn.Dropout(dropout * 0.7),
+                nn.Linear(input_dim // 2, num_classes)
+            )
+            self.feature_extractor = None
+            self.arcface = None
+
+        self.criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing)
+
+    def encode_tree(self, batch):
+        """Encode tree data"""
+        tree_feats = self.tree_encoder(batch)
+
+        if self.use_projection:
+            tree_embed = self.tree_projection(tree_feats)
+            return tree_embed
+        else:
+            return tree_feats
+
+    def encode_image(self, images):
+        """Encode image data"""
+        image_feats = self.image_encoder(images)
+
+        if self.use_projection:
+            image_embed = self.image_projection(image_feats)
+            return image_embed
+        else:
+            return image_feats
+
+    def forward(self, batch, return_features=False):
+        """
+        Forward pass
+
+        Args:
+            batch: contains batch.graph, batch.feats, batch.images, batch.label
+            return_features: if True, return embeddings along with logits
+        Returns:
+            loss: classification loss
+            logits: (B, num_classes)
+            features: (optional) embeddings
+        """
+        images = batch.images.cuda() if not batch.images.is_cuda else batch.images
+        labels = batch.label.cuda() if not batch.label.is_cuda else batch.label
+
+        if self.mode == 'tree_only':
+            tree_embed = self.encode_tree(batch)
+            tree_embed = F.normalize(tree_embed, dim=-1)
+            features = tree_embed
+
+        elif self.mode == 'image_only':
+            image_embed = self.encode_image(images)
+            image_embed = F.normalize(image_embed, dim=-1)
+            features = image_embed
+
+        else:  # multimodal
+            tree_embed = self.encode_tree(batch)
+            image_embed = self.encode_image(images)
+
+            tree_embed = F.normalize(tree_embed, dim=-1)
+            image_embed = F.normalize(image_embed, dim=-1)
+
+            if self.fusion_mode == 'concat':
+                features = torch.cat([tree_embed, image_embed], dim=1)
+            elif self.fusion_mode == 'add':
+                if hasattr(self, 'tree_dim_match'):
+                    tree_embed = self.tree_dim_match(tree_embed)
+                    image_embed = self.image_dim_match(image_embed)
+                features = tree_embed + image_embed
+            elif self.fusion_mode in ['gated', 'cmf', 'cross_attention', 'bi_attention', 'mhcma']:
+                if hasattr(self, 'tree_dim_match'):
+                    tree_embed = self.tree_dim_match(tree_embed)
+                    image_embed = self.image_dim_match(image_embed)
+                features = self.fusion_layer(tree_embed, image_embed)
+            else:
+                raise ValueError(f"Unknown fusion_mode: {self.fusion_mode}")
+
+        # Classification
+        if self.use_arcface:
+            extracted_features = self.feature_extractor(features)
+            logits = self.arcface(extracted_features, labels)
+            loss = self.criterion(logits, labels)
+        else:
+            logits = self.classifier(features)
+            loss = self.criterion(logits, labels)
+
+        if return_features:
+            return loss, logits, features
+        return loss, logits
+
+    def unfreeze_encoders(self):
+        """Unfreeze encoder weights for full fine-tuning"""
+        if self.mode in ['tree_only', 'multimodal'] and hasattr(self, 'tree_encoder'):
+            for param in self.tree_encoder.parameters():
+                param.requires_grad = True
+            if self.use_projection and hasattr(self, 'tree_projection'):
+                for param in self.tree_projection.parameters():
+                    param.requires_grad = True
+
+        if self.mode in ['image_only', 'multimodal'] and hasattr(self, 'image_encoder'):
+            for param in self.image_encoder.parameters():
+                param.requires_grad = True
+            if self.use_projection and hasattr(self, 'image_projection'):
+                for param in self.image_projection.parameters():
+                    param.requires_grad = True

graphformer/models/fusion.py

@@ -0,0 +1,339 @@
+"""
+Fusion Mechanisms for Multimodal Learning
+
+This module contains various fusion strategies to combine tree and image features:
+- CrossAttentionFusion: Cross-modal attention
+- CMF: Cross-Modal Fusion
+- BiDirectionalCrossAttention: Bidirectional cross-attention
+- GatedFusion: Gated fusion with learnable gates
+- MultiHeadCrossModalAttention: Multi-head cross-modal attention
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CrossAttentionFusion(nn.Module):
+    """Cross-attention fusion for tree and image features"""
+    def __init__(self, dim, num_heads=4, dropout=0.1):
+        super(CrossAttentionFusion, self).__init__()
+        self.num_heads = num_heads
+        self.dim = dim
+        self.head_dim = dim // num_heads
+
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+
+        # Query, Key, Value projections
+        self.q_proj = nn.Linear(dim, dim)
+        self.k_proj = nn.Linear(dim, dim)
+        self.v_proj = nn.Linear(dim, dim)
+        self.out_proj = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim),
+            nn.Dropout(dropout)
+        )
+
+        self.dropout = nn.Dropout(dropout)
+        self.scale = self.head_dim ** -0.5
+
+    def forward(self, tree_feat, image_feat):
+        """
+        Args:
+            tree_feat: (B, dim) tree features
+            image_feat: (B, dim) image features
+        Returns:
+            fused: (B, dim) fused features
+        """
+        B = tree_feat.shape[0]
+
+        # Add sequence dimension: (B, 1, dim)
+        tree_feat = tree_feat.unsqueeze(1)
+        image_feat = image_feat.unsqueeze(1)
+
+        # Tree attends to image (tree as query, image as key/value)
+        Q = self.q_proj(tree_feat).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        K = self.k_proj(image_feat).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        V = self.v_proj(image_feat).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Attention scores
+        attn = (Q @ K.transpose(-2, -1)) * self.scale
+        attn = torch.softmax(attn, dim=-1)
+        attn = self.dropout(attn)
+
+        # Apply attention to values
+        out = (attn @ V).transpose(1, 2).contiguous().view(B, 1, self.dim)
+
+        # Project and add residual
+        out = self.out_proj(out.squeeze(1))
+        fused = out + tree_feat.squeeze(1)
+
+        return fused
+
+
+class CMF(nn.Module):
+    """Cross-Modal Fusion with attention mechanism"""
+    def __init__(self, dim, dropout=0.1):
+        super(CMF, self).__init__()
+        self.dim = dim
+
+        # Feature-level attention
+        self.tree_attn = nn.Sequential(
+            nn.Linear(dim, dim),
+            nn.Tanh(),
+            nn.Linear(dim, 1)
+        )
+        self.image_attn = nn.Sequential(
+            nn.Linear(dim, dim),
+            nn.Tanh(),
+            nn.Linear(dim, 1)
+        )
+
+        # Cross-modal interaction
+        self.cross_proj = nn.Sequential(
+            nn.Linear(dim * 2, dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim, dim)
+        )
+
+    def forward(self, tree_feat, image_feat):
+        """
+        Args:
+            tree_feat: (B, dim)
+            image_feat: (B, dim)
+        Returns:
+            fused: (B, dim)
+        """
+        # Compute attention weights
+        tree_weight = torch.sigmoid(self.tree_attn(tree_feat))
+        image_weight = torch.sigmoid(self.image_attn(image_feat))
+
+        # Normalize weights
+        total_weight = tree_weight + image_weight + 1e-8
+        tree_weight = tree_weight / total_weight
+        image_weight = image_weight / total_weight
+
+        # Weighted combination
+        weighted_tree = tree_feat * tree_weight
+        weighted_image = image_feat * image_weight
+
+        # Cross-modal projection
+        combined = torch.cat([weighted_tree, weighted_image], dim=1)
+        fused = self.cross_proj(combined)
+
+        return fused
+
+
+class BiDirectionalCrossAttention(nn.Module):
+    """Bidirectional cross-attention: tree→image and image→tree"""
+    def __init__(self, dim, num_heads=4, dropout=0.1):
+        super(BiDirectionalCrossAttention, self).__init__()
+        self.num_heads = num_heads
+        self.dim = dim
+        self.head_dim = dim // num_heads
+
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+
+        # Tree → Image attention
+        self.tree2img_q = nn.Linear(dim, dim)
+        self.tree2img_k = nn.Linear(dim, dim)
+        self.tree2img_v = nn.Linear(dim, dim)
+        self.tree2img_out = nn.Linear(dim, dim)
+
+        # Image → Tree attention
+        self.img2tree_q = nn.Linear(dim, dim)
+        self.img2tree_k = nn.Linear(dim, dim)
+        self.img2tree_v = nn.Linear(dim, dim)
+        self.img2tree_out = nn.Linear(dim, dim)
+
+        self.dropout = nn.Dropout(dropout)
+        self.scale = self.head_dim ** -0.5
+
+        # Layer norm
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+    def _compute_attention(self, q_proj, k_proj, v_proj, query, key_value):
+        """Helper function to compute cross-attention"""
+        B = query.shape[0]
+
+        # Add sequence dimension
+        query = query.unsqueeze(1)  # (B, 1, dim)
+        key_value = key_value.unsqueeze(1)  # (B, 1, dim)
+
+        # Project
+        Q = q_proj(query).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        K = k_proj(key_value).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        V = v_proj(key_value).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Attention
+        attn = (Q @ K.transpose(-2, -1)) * self.scale
+        attn = torch.softmax(attn, dim=-1)
+        attn = self.dropout(attn)
+
+        # Apply to values
+        out = (attn @ V).transpose(1, 2).contiguous().view(B, 1, self.dim)
+        return out.squeeze(1)
+
+    def forward(self, tree_feat, image_feat):
+        """
+        Args:
+            tree_feat: (B, dim)
+            image_feat: (B, dim)
+        Returns:
+            tree_enhanced: (B, dim)
+            image_enhanced: (B, dim)
+        """
+        # Tree attends to image
+        tree_enhanced = self._compute_attention(
+            self.tree2img_q, self.tree2img_k, self.tree2img_v,
+            tree_feat, image_feat
+        )
+        tree_enhanced = self.tree2img_out(tree_enhanced)
+        tree_enhanced = self.norm1(tree_feat + tree_enhanced)
+
+        # Image attends to tree
+        image_enhanced = self._compute_attention(
+            self.img2tree_q, self.img2tree_k, self.img2tree_v,
+            image_feat, tree_feat
+        )
+        image_enhanced = self.img2tree_out(image_enhanced)
+        image_enhanced = self.norm2(image_feat + image_enhanced)
+
+        # Concatenate both enhanced features
+        fused = torch.cat([tree_enhanced, image_enhanced], dim=1)
+
+        return fused
+
+
+class GatedFusion(nn.Module):
+    """Gated fusion with learnable gates for tree and image modalities"""
+    def __init__(self, dim, dropout=0.1):
+        super(GatedFusion, self).__init__()
+        self.dim = dim
+
+        # Gating mechanism
+        self.gate_tree = nn.Sequential(
+            nn.Linear(dim * 2, dim),
+            nn.Sigmoid()
+        )
+        self.gate_image = nn.Sequential(
+            nn.Linear(dim * 2, dim),
+            nn.Sigmoid()
+        )
+
+        # Feature transformation
+        self.tree_transform = nn.Sequential(
+            nn.Linear(dim, dim),
+            nn.ReLU(),
+            nn.Dropout(dropout)
+        )
+        self.image_transform = nn.Sequential(
+            nn.Linear(dim, dim),
+            nn.ReLU(),
+            nn.Dropout(dropout)
+        )
+
+        # Output projection
+        self.output = nn.Linear(dim, dim)
+
+    def forward(self, tree_feat, image_feat):
+        """
+        Args:
+            tree_feat: (B, dim)
+            image_feat: (B, dim)
+        Returns:
+            fused: (B, dim)
+        """
+        # Concatenate features for gating
+        combined = torch.cat([tree_feat, image_feat], dim=1)
+
+        # Compute gates
+        gate_t = self.gate_tree(combined)
+        gate_i = self.gate_image(combined)
+
+        # Transform features
+        tree_transformed = self.tree_transform(tree_feat)
+        image_transformed = self.image_transform(image_feat)
+
+        # Apply gates
+        gated_tree = gate_t * tree_transformed
+        gated_image = gate_i * image_transformed
+
+        # Combine
+        fused = gated_tree + gated_image
+        fused = self.output(fused)
+
+        return fused
+
+
+class MultiHeadCrossModalAttention(nn.Module):
+    """Multi-head cross-modal attention for flexible fusion"""
+    def __init__(self, dim, num_heads=8, dropout=0.1):
+        super(MultiHeadCrossModalAttention, self).__init__()
+        self.num_heads = num_heads
+        self.dim = dim
+        self.head_dim = dim // num_heads
+
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+
+        # Projections
+        self.q_tree = nn.Linear(dim, dim)
+        self.k_tree = nn.Linear(dim, dim)
+        self.v_tree = nn.Linear(dim, dim)
+
+        self.q_image = nn.Linear(dim, dim)
+        self.k_image = nn.Linear(dim, dim)
+        self.v_image = nn.Linear(dim, dim)
+
+        self.out_proj = nn.Linear(dim * 2, dim)
+        self.dropout = nn.Dropout(dropout)
+        self.scale = self.head_dim ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, tree_feat, image_feat):
+        """
+        Args:
+            tree_feat: (B, dim)
+            image_feat: (B, dim)
+        Returns:
+            fused: (B, dim)
+        """
+        B = tree_feat.shape[0]
+
+        # Add sequence dimension
+        tree_feat = tree_feat.unsqueeze(1)  # (B, 1, dim)
+        image_feat = image_feat.unsqueeze(1)  # (B, 1, dim)
+
+        # Project tree features
+        Q_t = self.q_tree(tree_feat).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        K_t = self.k_tree(tree_feat).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        V_t = self.v_tree(tree_feat).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Project image features
+        Q_i = self.q_image(image_feat).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        K_i = self.k_image(image_feat).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        V_i = self.v_image(image_feat).view(B, 1, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Tree self-attention with image context
+        attn_t = (Q_t @ K_i.transpose(-2, -1)) * self.scale
+        attn_t = torch.softmax(attn_t, dim=-1)
+        attn_t = self.dropout(attn_t)
+        out_t = (attn_t @ V_i).transpose(1, 2).contiguous().view(B, 1, self.dim)
+
+        # Image self-attention with tree context
+        attn_i = (Q_i @ K_t.transpose(-2, -1)) * self.scale
+        attn_i = torch.softmax(attn_i, dim=-1)
+        attn_i = self.dropout(attn_i)
+        out_i = (attn_i @ V_t).transpose(1, 2).contiguous().view(B, 1, self.dim)
+
+        # Concatenate and project
+        combined = torch.cat([out_t.squeeze(1), out_i.squeeze(1)], dim=1)
+        fused = self.out_proj(combined)
+        fused = self.norm(fused)
+
+        return fused

graphformer/models/hybrid_resnet_persistencevit.py

@@ -0,0 +1,414 @@
+"""
+Hybrid ResNet-PersistenceViT Architecture
+
+Combines:
+- ResNet's early convolutional layers (conv1, bn1, relu, maxpool, layer1, layer2)
+- PersistenceViT's topological attention mechanism (replacing layer3, layer4)
+
+This hybrid leverages:
+1. ResNet's proven feature extraction in lower layers
+2. PersistenceViT's topological awareness in higher layers
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torchvision.models import ResNet18_Weights, ResNet50_Weights
+
+
+# Import PersistenceViT components
+from .image_encoder import (
+    BirthDeathAttention,
+    BirthDeathTransformerBlock,
+    TopologicalFeatureAggregation,
+)
+
+
+class ResNetFeatureExtractor(nn.Module):
+    """Extract features from early ResNet layers"""
+    def __init__(self, model_type='resnet18', num_layers=2):
+        """
+        Args:
+            model_type: 'resnet18' or 'resnet50'
+            num_layers: How many ResNet layers to keep (1, 2, or 3)
+                       1 = conv1 + layer1 (64 channels, /4 spatial)
+                       2 = conv1 + layer1 + layer2 (128/512 channels, /8 spatial)
+                       3 = conv1 + layer1 + layer2 + layer3 (256/1024 channels, /16 spatial)
+        """
+        super().__init__()
+
+        if model_type == 'resnet18':
+            resnet = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
+            self.layer_dims = [64, 128, 256, 512]
+        elif model_type == 'resnet50':
+            resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+            self.layer_dims = [256, 512, 1024, 2048]
+        else:
+            raise ValueError(f"model_type must be 'resnet18' or 'resnet50'")
+
+        self.model_type = model_type
+        self.num_layers = num_layers
+
+        # Initial conv layers
+        self.conv1 = resnet.conv1
+        self.bn1 = resnet.bn1
+        self.relu = resnet.relu
+        self.maxpool = resnet.maxpool
+
+        # ResNet layers
+        self.layer1 = resnet.layer1
+        if num_layers >= 2:
+            self.layer2 = resnet.layer2
+        if num_layers >= 3:
+            self.layer3 = resnet.layer3
+
+        # Output channels
+        self.out_channels = self.layer_dims[num_layers - 1]
+
+    def forward(self, x):
+        """
+        Args:
+            x: (B, 3, H, W)
+        Returns:
+            features: (B, out_channels, H', W') where H' = H / (4 * 2^num_layers)
+        """
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        if self.num_layers >= 2:
+            x = self.layer2(x)
+        if self.num_layers >= 3:
+            x = self.layer3(x)
+
+        return x
+
+
+class CNNToTransformerAdapter(nn.Module):
+    """Convert CNN feature maps to transformer tokens"""
+    def __init__(self, in_channels, dim, patch_size=2):
+        """
+        Args:
+            in_channels: Number of input channels from CNN
+            dim: Transformer hidden dimension
+            patch_size: Size of patches to group CNN features (default: 2x2)
+        """
+        super().__init__()
+        self.patch_size = patch_size
+
+        # 1x1 conv to reduce channels, then adaptive pooling
+        self.proj = nn.Sequential(
+            nn.Conv2d(in_channels, dim, kernel_size=patch_size, stride=patch_size),
+            nn.BatchNorm2d(dim),
+            nn.ReLU(inplace=True)
+        )
+
+        # Importance estimator (similar to PersistenceViT)
+        self.importance_estimator = nn.Sequential(
+            nn.Conv2d(in_channels, 32, kernel_size=patch_size, stride=patch_size),
+            nn.ReLU(),
+            nn.Conv2d(32, 1, kernel_size=1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        """
+        Args:
+            x: (B, in_channels, H, W)
+        Returns:
+            tokens: (B, N, dim) where N = (H/patch_size) * (W/patch_size)
+            importance: (B, N)
+        """
+        B, C, H, W = x.shape
+
+        # Project to embedding dimension
+        tokens = self.proj(x)  # (B, dim, H', W')
+        tokens = tokens.flatten(2).transpose(1, 2)  # (B, N, dim)
+
+        # Compute importance weights
+        importance = self.importance_estimator(x)  # (B, 1, H', W')
+        importance = importance.flatten(2).transpose(1, 2).squeeze(-1)  # (B, N)
+
+        # Weight tokens by importance
+        weighted_tokens = tokens * (1 + importance.unsqueeze(-1))
+
+        return weighted_tokens, importance
+
+
+class HybridResNetPersistenceViT(nn.Module):
+    """
+    Hybrid architecture combining ResNet and PersistenceViT
+
+    Architecture:
+    1. ResNet early layers (conv1 + layer1 + layer2) - proven feature extraction
+    2. CNN-to-Transformer adapter - convert feature maps to tokens
+    3. PersistenceViT attention blocks - topological reasoning
+    4. Classification head
+    """
+    def __init__(
+        self,
+        output_dim=128,
+        image_size=224,
+        resnet_type='resnet18',
+        resnet_layers=2,
+        dim=256,
+        depth=4,
+        heads=8,
+        mlp_dim=512,
+        dropout=0.2,
+        homology_dims=3,
+        freeze_resnet=False,
+    ):
+        """
+        Args:
+            output_dim: Output feature dimension
+            image_size: Input image size
+            resnet_type: 'resnet18' or 'resnet50'
+            resnet_layers: Number of ResNet layers to keep (1, 2, or 3)
+            dim: Transformer hidden dimension
+            depth: Number of transformer blocks
+            heads: Number of attention heads
+            mlp_dim: MLP hidden dimension
+            dropout: Dropout rate
+            homology_dims: Number of homology dimension tokens
+            freeze_resnet: Freeze ResNet layers
+        """
+        super().__init__()
+
+        self.image_size = image_size
+        self.dim = dim
+
+        # 1. ResNet feature extractor
+        self.resnet_extractor = ResNetFeatureExtractor(
+            model_type=resnet_type,
+            num_layers=resnet_layers
+        )
+
+        if freeze_resnet:
+            for param in self.resnet_extractor.parameters():
+                param.requires_grad = False
+
+        # 2. CNN-to-Transformer adapter
+        # After resnet_layers=2: spatial size is image_size / 8
+        # We further reduce by patch_size=2, so final is image_size / 16
+        self.adapter = CNNToTransformerAdapter(
+            in_channels=self.resnet_extractor.out_channels,
+            dim=dim,
+            patch_size=2
+        )
+
+        # Calculate number of patches
+        spatial_reduction = 8 * 2  # ResNet (layer2) + adapter patch_size
+        self.num_patches = (image_size // spatial_reduction) ** 2
+
+        # 3. Special tokens (like PersistenceViT)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.homology_tokens = nn.Parameter(torch.randn(1, homology_dims, dim))
+
+        # 4. Positional encoding
+        num_special_tokens = 1 + homology_dims
+        self.pos_embedding = nn.Parameter(
+            torch.randn(1, num_special_tokens + self.num_patches, dim)
+        )
+        self.dropout = nn.Dropout(dropout)
+
+        # 5. PersistenceViT transformer blocks
+        self.transformer_blocks = nn.ModuleList([
+            BirthDeathTransformerBlock(dim, heads, mlp_dim, dropout)
+            for _ in range(depth)
+        ])
+
+        # 6. Topological aggregation
+        self.topological_aggregation = TopologicalFeatureAggregation(
+            dim, homology_dims
+        )
+
+        # 7. Output head
+        self.norm = nn.LayerNorm(dim)
+        self.mlp_head = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(dim, output_dim)
+        )
+
+    def forward(self, images):
+        """
+        Args:
+            images: (B, 3, H, W)
+        Returns:
+            features: (B, output_dim)
+        """
+        B = images.shape[0]
+
+        # 1. Extract features with ResNet
+        cnn_features = self.resnet_extractor(images)  # (B, C, H', W')
+
+        # 2. Convert to transformer tokens
+        tokens, importance_weights = self.adapter(cnn_features)  # (B, N, dim)
+
+        # 3. Add special tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        homology_tokens = self.homology_tokens.expand(B, -1, -1)
+        x = torch.cat([cls_tokens, homology_tokens, tokens], dim=1)  # (B, 1+H+N, dim)
+
+        # 4. Add positional encoding
+        x = x + self.pos_embedding[:, :(x.shape[1])]
+        x = self.dropout(x)
+
+        # 5. Pad importance weights for special tokens
+        num_special_tokens = 1 + self.homology_tokens.shape[1]
+        importance_padding = torch.zeros(B, num_special_tokens, device=importance_weights.device)
+        importance_weights_padded = torch.cat([importance_padding, importance_weights], dim=1)
+
+        # 6. Apply transformer blocks with topological attention
+        for block in self.transformer_blocks:
+            x = block(x, importance_weights_padded)
+
+        x = self.norm(x)
+
+        # 7. Topological aggregation (use cls + homology tokens)
+        x = self.topological_aggregation(x[:, :num_special_tokens])
+
+        # 8. Output projection
+        output = self.mlp_head(x)
+
+        return output
+
+
+class HybridImageEncoder(nn.Module):
+    """
+    Wrapper to integrate HybridResNetPersistenceViT into existing codebase
+    Compatible with ImageEncoder interface
+    """
+    def __init__(
+        self,
+        output_dim=128,
+        image_size=224,
+        resnet_type='resnet18',
+        resnet_layers=2,
+        dim=256,
+        depth=4,
+        heads=8,
+        mlp_dim=512,
+        dropout=0.2,
+        homology_dims=3,
+        freeze_resnet=False,
+    ):
+        super().__init__()
+
+        self.encoder = HybridResNetPersistenceViT(
+            output_dim=output_dim,
+            image_size=image_size,
+            resnet_type=resnet_type,
+            resnet_layers=resnet_layers,
+            dim=dim,
+            depth=depth,
+            heads=heads,
+            mlp_dim=mlp_dim,
+            dropout=dropout,
+            homology_dims=homology_dims,
+            freeze_resnet=freeze_resnet,
+        )
+
+    def forward(self, images, persistence_coords=None, pixel_coords=None):
+        """
+        Args:
+            images: (B, 3, H, W)
+            persistence_coords: ignored (for interface compatibility)
+            pixel_coords: ignored (for interface compatibility)
+        Returns:
+            features: (B, output_dim)
+        """
+        return self.encoder(images)
+
+
+# ============================================================================
+# Integration into existing ImageEncoder
+# ============================================================================
+
+def create_hybrid_encoder(output_dim=128, image_size=224, freeze_backbone=False, **kwargs):
+    """
+    Factory function to create hybrid encoder
+
+    Args:
+        output_dim: Output dimension
+        image_size: Input image size
+        freeze_backbone: Freeze ResNet layers
+        **kwargs: Additional arguments
+            - resnet_type: 'resnet18' or 'resnet50' (default: 'resnet50')
+            - resnet_layers: Number of ResNet layers (1, 2, or 3) (default: 2)
+            - dim: Transformer dimension (default: 256)
+            - depth: Number of transformer blocks (default: 4)
+            - heads: Number of attention heads (default: 8)
+    """
+    resnet_type = kwargs.get('resnet_type', 'resnet50')
+    resnet_layers = kwargs.get('resnet_layers', 2)
+    dim = kwargs.get('dim', 256)
+    depth = kwargs.get('depth', 4)
+    heads = kwargs.get('heads', 8)
+    mlp_dim = kwargs.get('mlp_dim', 512)
+    dropout = kwargs.get('dropout', 0.2)
+    homology_dims = kwargs.get('homology_dims', 3)
+
+    return HybridImageEncoder(
+        output_dim=output_dim,
+        image_size=image_size,
+        resnet_type=resnet_type,
+        resnet_layers=resnet_layers,
+        dim=dim,
+        depth=depth,
+        heads=heads,
+        mlp_dim=mlp_dim,
+        dropout=dropout,
+        homology_dims=homology_dims,
+        freeze_resnet=freeze_backbone,
+    )
+
+
+if __name__ == "__main__":
+    # Test the hybrid model
+    print("Testing HybridResNetPersistenceViT...")
+
+    # Create model
+    model = HybridResNetPersistenceViT(
+        output_dim=256,
+        image_size=224,
+        resnet_type='resnet18',
+        resnet_layers=2,
+        dim=256,
+        depth=4,
+        heads=8,
+        mlp_dim=512,
+        dropout=0.2,
+        homology_dims=3,
+        freeze_resnet=False,
+    )
+
+    # Count parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print(f"Total parameters: {total_params:,}")
+    print(f"Trainable parameters: {trainable_params:,}")
+
+    # Test forward pass
+    batch_size = 4
+    images = torch.randn(batch_size, 3, 224, 224)
+
+    output = model(images)
+    print(f"Input shape: {images.shape}")
+    print(f"Output shape: {output.shape}")
+
+    # Component breakdown
+    resnet_params = sum(p.numel() for p in model.resnet_extractor.parameters())
+    adapter_params = sum(p.numel() for p in model.adapter.parameters())
+    transformer_params = sum(p.numel() for p in model.transformer_blocks.parameters())
+
+    print(f"\nComponent breakdown:")
+    print(f"  ResNet layers: {resnet_params:,}")
+    print(f"  CNN-to-Transformer adapter: {adapter_params:,}")
+    print(f"  Transformer blocks: {transformer_params:,}")
+
+    print("\nHybrid model test passed!")

graphformer/models/image_encoder.py

@@ -0,0 +1,702 @@
+"""
+Image Encoder Components for Persistence Images
+
+This module contains various image encoders optimized for persistence images:
+- SimpleCNN: Lightweight CNN
+- SmallViT: Compact Vision Transformer
+- PersistenceViT: Topologically-aware Vision Transformer with persistence-weighted positional encoding
+- ResNet18/ResNet50: Standard ResNet encoders
+- DINOv2: Self-supervised visual encoders
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torchvision.models import ResNet18_Weights, ResNet50_Weights, ResNet101_Weights
+from typing import Tuple, Optional
+
+
+# ============================================================================
+# Simple CNN Encoder
+# ============================================================================
+
+class SimpleCNN(nn.Module):
+    """Lightweight CNN for persistence images"""
+    def __init__(self, output_dim=128):
+        super(SimpleCNN, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(3, 32, 3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, 3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Dropout2d(0.3),
+            nn.Conv2d(64, 128, 3, padding=1),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d(1)
+        )
+        self.fc = nn.Linear(128, output_dim)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+
+# ============================================================================
+# Vision Transformer Components
+# ============================================================================
+
+class MultiHeadAttention(nn.Module):
+    """Multi-head self-attention for Vision Transformer"""
+    def __init__(self, dim, heads, dropout=0.1):
+        super().__init__()
+        self.heads = heads
+        self.scale = (dim // heads) ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.attn_drop = nn.Dropout(dropout)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(dropout)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.heads, C // self.heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    """Transformer block with self-attention and MLP"""
+    def __init__(self, dim, heads, mlp_dim, dropout=0.1):
+        super().__init__()
+        self.attention = MultiHeadAttention(dim, heads, dropout)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, mlp_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(mlp_dim, dim),
+            nn.Dropout(dropout)
+        )
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        x = x + self.attention(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class SmallViT(nn.Module):
+    """Lightweight Vision Transformer for persistence images"""
+    def __init__(
+        self,
+        image_size=224,
+        patch_size=14,
+        output_dim=128,
+        dim=128,
+        depth=4,
+        heads=4,
+        mlp_dim=256,
+        channels=3,
+        dropout=0.2,
+        emb_dropout=0.2
+    ):
+        super().__init__()
+        assert image_size % patch_size == 0
+        num_patches = (image_size // patch_size) ** 2
+
+        self.patch_embedding = nn.Sequential(
+            nn.Conv2d(channels, dim, kernel_size=patch_size, stride=patch_size),
+            nn.Flatten(2),
+        )
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = nn.ModuleList([
+            TransformerBlock(dim, heads, mlp_dim, dropout)
+            for _ in range(depth)
+        ])
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Dropout(dropout),
+            nn.Linear(dim, output_dim)
+        )
+
+    def forward(self, img):
+        x = self.patch_embedding(img).transpose(1, 2)
+        b, n, _ = x.shape
+        cls_tokens = self.cls_token.expand(b, -1, -1)
+        x = torch.cat([cls_tokens, x], dim=1)
+        x = x + self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        for transformer_block in self.transformer:
+            x = transformer_block(x)
+
+        return self.mlp_head(x[:, 0])
+
+
+# ============================================================================
+# Topologically-Aware Vision Transformer (PersistenceViT)
+# ============================================================================
+
+class TopologicalPatchEmbedding(nn.Module):
+    """Topological-Aware Patch Embedding with importance weighting"""
+    def __init__(self, in_channels, dim, patch_size, image_size):
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_patches = (image_size // patch_size) ** 2
+
+        self.proj = nn.Conv2d(in_channels, dim, kernel_size=patch_size, stride=patch_size)
+
+        # Importance estimator for high-persistence regions
+        self.importance_estimator = nn.Sequential(
+            nn.Conv2d(in_channels, 32, kernel_size=patch_size, stride=patch_size),
+            nn.ReLU(),
+            nn.Conv2d(32, 1, kernel_size=1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        patches = self.proj(x).flatten(2).transpose(1, 2)  # (B, N, dim)
+        importance = self.importance_estimator(x).flatten(2).transpose(1, 2)  # (B, N, 1)
+        weighted_patches = patches * (1 + importance)
+        return weighted_patches, importance.squeeze(-1)
+
+
+class BirthDeathAttention(nn.Module):
+    """Attention mechanism with persistence-based modulation"""
+    def __init__(self, dim, heads, dropout=0.1):
+        super().__init__()
+        self.heads = heads
+        self.scale = (dim // heads) ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.attn_drop = nn.Dropout(dropout)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(dropout)
+        self.persistence_bias = nn.Parameter(torch.zeros(1, heads, 1, 1))
+
+    def forward(self, x, importance_weights=None):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.heads, C // self.heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+
+        if importance_weights is not None:
+            importance = importance_weights.unsqueeze(1).unsqueeze(-1)
+            attn = attn + self.persistence_bias + importance * 0.1
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class BirthDeathTransformerBlock(nn.Module):
+    """Transformer block with topological awareness"""
+    def __init__(self, dim, heads, mlp_dim, dropout=0.1):
+        super().__init__()
+        self.attention = BirthDeathAttention(dim, heads, dropout)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, mlp_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(mlp_dim, dim),
+            nn.Dropout(dropout)
+        )
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+    def forward(self, x, importance_weights=None):
+        x = x + self.attention(self.norm1(x), importance_weights)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class TopologicalFeatureAggregation(nn.Module):
+    """Aggregates CLS token with homology-specific tokens"""
+    def __init__(self, dim, num_homology_dims):
+        super().__init__()
+        self.num_homology = num_homology_dims
+        self.cross_attn = nn.MultiheadAttention(dim, num_heads=4, dropout=0.1, batch_first=True)
+        self.norm = nn.LayerNorm(dim)
+        self.aggregation_weights = nn.Parameter(torch.ones(1, num_homology_dims + 1, 1))
+
+    def forward(self, tokens):
+        B = tokens.shape[0]
+        cls_token = tokens[:, 0:1]
+        homology_tokens = tokens[:, 1:]
+
+        aggregated, _ = self.cross_attn(cls_token, homology_tokens, homology_tokens)
+        aggregated = self.norm(aggregated)
+
+        weights = F.softmax(self.aggregation_weights, dim=1)
+        weighted_tokens = tokens * weights
+        final = aggregated + weighted_tokens.sum(dim=1, keepdim=True)
+
+        return final.squeeze(1)
+
+
+class PersistenceWeightedPositionalEncoding(nn.Module):
+    """
+    NOVEL CONTRIBUTION: Positional encoding based on topological persistence
+    rather than just spatial grid location.
+
+    Key insight: Patches corresponding to high-persistence features should have
+    similar positional encodings regardless of their spatial location, because
+    they represent similar topological significance.
+    """
+    def __init__(self, dim, image_size=224, patch_size=14):
+        super().__init__()
+        self.dim = dim
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_patches = (image_size // patch_size) ** 2
+
+        # Spatial positional encoding (standard grid-based)
+        self.spatial_pos = nn.Parameter(torch.randn(1, self.num_patches, dim // 2))
+
+        # Persistence-based encoding
+        self.birth_encoder = nn.Sequential(
+            nn.Linear(1, dim // 4),
+            nn.ReLU(),
+            nn.Linear(dim // 4, dim // 4)
+        )
+
+        self.persistence_encoder = nn.Sequential(
+            nn.Linear(1, dim // 4),
+            nn.ReLU(),
+            nn.Linear(dim // 4, dim // 4)
+        )
+
+        # Fusion of spatial and topological encodings
+        self.fusion = nn.Sequential(
+            nn.Linear(dim, dim),
+            nn.LayerNorm(dim),
+            nn.Tanh()
+        )
+
+    def encode_patch_persistence(
+        self,
+        persistence_coords: torch.Tensor,  # (N_features, 3)
+        pixel_coords: torch.Tensor         # (N_features, 2)
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Assign persistence values to patches based on which features fall in each patch.
+
+        Returns:
+            patch_birth: (num_patches, 1) - average birth for each patch
+            patch_persistence: (num_patches, 1) - average persistence for each patch
+        """
+        device = persistence_coords.device
+        num_patches_h = self.image_size // self.patch_size
+        num_patches_w = self.image_size // self.patch_size
+
+        # Initialize patch statistics
+        patch_birth_sum = torch.zeros(num_patches_h, num_patches_w, device=device)
+        patch_pers_sum = torch.zeros(num_patches_h, num_patches_w, device=device)
+        patch_count = torch.zeros(num_patches_h, num_patches_w, device=device)
+
+        # Assign features to patches
+        for i in range(len(pixel_coords)):
+            # Skip padded entries (all zeros)
+            if torch.all(pixel_coords[i] == 0) and torch.all(persistence_coords[i] == 0):
+                continue
+
+            x, y = pixel_coords[i]
+            patch_x = int(torch.clamp(x / self.patch_size, 0, num_patches_w - 1))
+            patch_y = int(torch.clamp(y / self.patch_size, 0, num_patches_h - 1))
+
+            birth = persistence_coords[i, 0]
+            pers = persistence_coords[i, 2]
+
+            patch_birth_sum[patch_y, patch_x] += birth
+            patch_pers_sum[patch_y, patch_x] += pers
+            patch_count[patch_y, patch_x] += 1
+
+        # Average persistence values per patch
+        mask = patch_count > 0
+        patch_birth_avg = torch.zeros_like(patch_birth_sum)
+        patch_pers_avg = torch.zeros_like(patch_pers_sum)
+
+        patch_birth_avg[mask] = patch_birth_sum[mask] / patch_count[mask]
+        patch_pers_avg[mask] = patch_pers_sum[mask] / patch_count[mask]
+
+        # Flatten to (num_patches, 1)
+        patch_birth_flat = patch_birth_avg.flatten().unsqueeze(-1)  # (num_patches, 1)
+        patch_pers_flat = patch_pers_avg.flatten().unsqueeze(-1)    # (num_patches, 1)
+
+        return patch_birth_flat, patch_pers_flat
+
+    def forward(
+        self,
+        batch_size: int,
+        persistence_coords: Optional[torch.Tensor] = None,
+        pixel_coords: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Generate positional encoding combining spatial and persistence information.
+
+        Args:
+            batch_size: Batch size
+            persistence_coords: (B, N_features, 3) - optional persistence coordinates
+            pixel_coords: (B, N_features, 2) - optional pixel coordinates
+
+        Returns:
+            pos_encoding: (B, num_patches, dim)
+        """
+        # Start with spatial encoding
+        spatial_enc = self.spatial_pos.expand(batch_size, -1, -1)  # (B, num_patches, dim//2)
+
+        if persistence_coords is None or pixel_coords is None:
+            # Fall back to spatial-only encoding
+            zero_enc = torch.zeros(batch_size, self.num_patches, self.dim // 2,
+                                  device=spatial_enc.device)
+            return torch.cat([spatial_enc, zero_enc], dim=-1)
+
+        # Encode persistence information for each sample in batch
+        batch_persistence_enc = []
+
+        for b in range(batch_size):
+            pers_coords_b = persistence_coords[b]  # (N_features, 3)
+            pix_coords_b = pixel_coords[b]         # (N_features, 2)
+
+            # Get per-patch persistence statistics
+            patch_birth, patch_pers = self.encode_patch_persistence(pers_coords_b, pix_coords_b)
+
+            # Encode through MLPs
+            birth_enc = self.birth_encoder(patch_birth)        # (num_patches, dim//4)
+            pers_enc = self.persistence_encoder(patch_pers)    # (num_patches, dim//4)
+
+            # Combine
+            persistence_enc = torch.cat([birth_enc, pers_enc], dim=-1)  # (num_patches, dim//2)
+            batch_persistence_enc.append(persistence_enc)
+
+        persistence_enc = torch.stack(batch_persistence_enc, dim=0)  # (B, num_patches, dim//2)
+
+        # Combine spatial and persistence encodings
+        combined = torch.cat([spatial_enc, persistence_enc], dim=-1)  # (B, num_patches, dim)
+
+        # Fuse through learned transformation
+        pos_encoding = self.fusion(combined)
+
+        return pos_encoding
+
+
+class PersistenceViT(nn.Module):
+    """
+    Vision Transformer for Persistence Images with topological inductive biases.
+
+    Features:
+    - Topological-aware patch embedding with importance weighting
+    - Birth-death attention mechanism
+    - Multi-scale persistence encoding via homology tokens
+    - Topological feature aggregation
+    - Persistence-weighted positional encoding
+    """
+    def __init__(
+        self,
+        image_size=256,
+        patch_size=16,
+        output_dim=128,
+        dim=128,
+        depth=4,
+        heads=4,
+        mlp_dim=256,
+        channels=3,
+        dropout=0.2,
+        homology_dims=3,
+    ):
+        super().__init__()
+
+        self.image_size = image_size
+        self.patch_size = patch_size
+        num_patches = (image_size // patch_size) ** 2
+
+        # Topological-aware patch embedding
+        self.patch_embedding = TopologicalPatchEmbedding(
+            channels, dim, patch_size, image_size
+        )
+
+        # Multi-scale persistence tokens
+        self.homology_tokens = nn.Parameter(torch.randn(1, homology_dims, dim))
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+
+        # Persistence-weighted positional encoding
+        self.pos_encoder = PersistenceWeightedPositionalEncoding(
+            dim=dim,
+            image_size=image_size,
+            patch_size=patch_size
+        )
+        self.dropout = nn.Dropout(dropout)
+
+        # Birth-death transformer blocks
+        self.transformer_blocks = nn.ModuleList([
+            BirthDeathTransformerBlock(dim, heads, mlp_dim, dropout)
+            for _ in range(depth)
+        ])
+
+        # Topological aggregation
+        self.topological_aggregation = TopologicalFeatureAggregation(dim, homology_dims)
+
+        # Output head
+        self.norm = nn.LayerNorm(dim)
+        self.mlp_head = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(dim, output_dim)
+        )
+
+    def forward(self, persistence_img, persistence_coords=None, pixel_coords=None):
+        """
+        Forward pass with optional persistence coordinate information.
+
+        Args:
+            persistence_img: (B, C, H, W) - persistence image
+            persistence_coords: (B, N_features, 3) - optional normalized (birth, death, persistence)
+            pixel_coords: (B, N_features, 2) - optional (x, y) pixel locations
+
+        Returns:
+            output: (B, output_dim) - encoded representation
+        """
+        B = persistence_img.shape[0]
+
+        # Topological-aware patch embedding
+        x, importance_weights = self.patch_embedding(persistence_img)  # (B, N, dim)
+
+        # Generate persistence-weighted positional encoding
+        pos_encoding = self.pos_encoder(
+            batch_size=B,
+            persistence_coords=persistence_coords,
+            pixel_coords=pixel_coords
+        )  # (B, num_patches, dim)
+
+        # Add positional encoding
+        x = x + pos_encoding
+
+        # Add special tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        homology_tokens = self.homology_tokens.expand(B, -1, -1)
+        x = torch.cat([cls_tokens, homology_tokens, x], dim=1)
+
+        # Pad importance_weights for special tokens (cls + homology)
+        num_special_tokens = 1 + self.homology_tokens.shape[1]
+        importance_padding = torch.zeros(B, num_special_tokens, device=importance_weights.device)
+        importance_weights_padded = torch.cat([importance_padding, importance_weights], dim=1)
+
+        x = self.dropout(x)
+
+        # Birth-death transformer blocks
+        for block in self.transformer_blocks:
+            x = block(x, importance_weights_padded)
+
+        x = self.norm(x)
+
+        # Topological aggregation
+        x = self.topological_aggregation(x[:, :1+self.homology_tokens.shape[1]])
+
+        return self.mlp_head(x)
+
+
+# ============================================================================
+# Unified Image Encoder Interface
+# ============================================================================
+
+class ImageEncoder(nn.Module):
+    """
+    Unified image encoder interface supporting multiple architectures:
+    - SimpleCNN: Lightweight CNN
+    - SmallViT: Compact Vision Transformer
+    - PersistenceViT: Topologically-aware ViT
+    - ResNet18/ResNet50/ResNet101: Standard CNNs with ImageNet pretraining
+    - DINOv2: Self-supervised vision encoders (ViT-S/B/L/g)
+    - ConvNeXt-Small: Modern CNN architecture
+    - HybridResNetViT: ResNet conv layers + PersistenceViT attention
+    """
+    def __init__(self, output_dim=128, model_type='resnet18', image_size=224, freeze_backbone=False):
+        super(ImageEncoder, self).__init__()
+
+        if model_type == 'hybrid_resnet18_vit':
+            from .hybrid_resnet_persistencevit import create_hybrid_encoder
+            self.encoder = create_hybrid_encoder(
+                output_dim=output_dim,
+                image_size=image_size,
+                freeze_backbone=freeze_backbone,
+                resnet_type='resnet18',
+                resnet_layers=2,
+                dim=256,
+                depth=4,
+                heads=8,
+                mlp_dim=512,
+                dropout=0.2,
+            )
+            self.use_simple_model = True
+
+        elif model_type == 'hybrid_resnet50_vit':
+            from .hybrid_resnet_persistencevit import create_hybrid_encoder
+            self.encoder = create_hybrid_encoder(
+                output_dim=output_dim,
+                image_size=image_size,
+                freeze_backbone=freeze_backbone,
+                resnet_type='resnet50',
+                resnet_layers=2,
+                dim=256,
+                depth=4,
+                heads=8,
+                mlp_dim=512,
+                dropout=0.2,
+            )
+            self.use_simple_model = True
+            
+        elif model_type == 'simplecnn':
+            self.encoder = SimpleCNN(output_dim=output_dim)
+            self.use_simple_model = True
+        elif model_type == 'smallvit':
+            self.encoder = SmallViT(image_size=image_size, patch_size=16, output_dim=output_dim,
+                                   dim=128, depth=6, heads=4, mlp_dim=256, channels=3,
+                                   dropout=0.2, emb_dropout=0.2)
+            self.use_simple_model = True
+        elif model_type == 'persistencevit':
+            self.encoder = PersistenceViT(image_size=image_size, patch_size=16, output_dim=output_dim,
+                                         dim=128, depth=6, heads=4, mlp_dim=256, channels=3,
+                                         dropout=0.2, homology_dims=3)
+            self.use_simple_model = True
+        elif model_type == 'resnet18':
+            resnet = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
+            feature_dim = 512
+            self.encoder = nn.Sequential(*list(resnet.children())[:-1])
+            self.fc = nn.Sequential(
+                 nn.Linear(feature_dim, 256),
+                nn.LayerNorm(256),
+                nn.ReLU(),
+                nn.Linear(256, output_dim),
+            )
+            self.use_simple_model = False
+        elif model_type == 'resnet50':
+            resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+            feature_dim = 2048
+            self.encoder = nn.Sequential(*list(resnet.children())[:-1])
+            self.fc = nn.Sequential(
+                nn.Linear(feature_dim, 256),
+                nn.LayerNorm(256),
+                nn.ReLU(),
+                nn.Linear(256, output_dim),
+            )
+            self.use_simple_model = False
+        elif model_type == 'resnet101':
+            resnet = models.resnet101(weights=ResNet101_Weights.IMAGENET1K_V1)
+            feature_dim = 2048
+            self.encoder = nn.Sequential(*list(resnet.children())[:-1])
+            self.fc = nn.Sequential(
+                nn.Linear(feature_dim, 256)
+            )
+            self.use_simple_model = False
+        elif model_type.startswith('dinov2'):
+            self.encoder = DINOv2ImageEncoder(output_dim=output_dim, freeze_backbone=False, model_variant=model_type)
+        
+            self.use_simple_model = True
+        elif model_type == 'convnext_small':
+            from torchvision.models import convnext_small, ConvNeXt_Small_Weights
+            convnext = convnext_small(weights=ConvNeXt_Small_Weights.IMAGENET1K_V1)
+            feature_dim = 768
+            self.encoder = nn.Sequential(*list(convnext.children())[:-1])
+            if freeze_backbone:
+                for param in self.encoder.parameters():
+                    param.requires_grad = False
+            self.fc = nn.Sequential(
+                nn.Flatten(),
+                nn.Linear(feature_dim, 256),
+                nn.LayerNorm(256),
+                nn.ReLU(),
+                nn.Linear(256, output_dim),
+            )
+            self.use_simple_model = False
+        else:
+            raise ValueError(f"model_type must be 'simplecnn', 'smallvit', 'persistencevit', 'resnet18', 'resnet50', 'resnet101', 'dinov2_vits14', 'dinov2_vitb14', 'dinov2_vitl14', 'dinov2_vitg14', 'convnext_small', or 'hybrid_resnet_vit', got {model_type}")
+
+        self.model_type = model_type
+
+    def forward(self, images, persistence_coords=None, pixel_coords=None):
+        """
+        Args:
+            images: (B, 3, H, W) or (B, 1, H, W) tensor of persistence images
+            persistence_coords: (B, N_features, 3) - optional for PersistenceViT
+            pixel_coords: (B, N_features, 2) - optional for PersistenceViT
+        Returns:
+            features: (B, output_dim) tensor of image features
+        """
+        if self.use_simple_model:
+            # PersistenceViT can use the coordinates
+            if self.model_type == 'persistencevit':
+                return self.encoder(images, persistence_coords, pixel_coords)
+            else:
+                return self.encoder(images)
+        else:
+            features = self.encoder(images)  # (B, feature_dim, 1, 1)
+            features = features.view(features.size(0), -1)  # (B, feature_dim)
+            features = self.fc(features)  # (B, output_dim)
+            return features
+
+
+class DINOv2ImageEncoder(nn.Module):
+    """
+    DINOv2-based image encoder for persistence images.
+    Supports multiple DINOv2 variants: ViT-S/14, ViT-B/14, ViT-L/14, ViT-g/14.
+    Returns raw features (384-dim for vits14) without projection head.
+    """
+    def __init__(self, output_dim=128, freeze_backbone=True, model_variant='dinov2_vits14'):
+        super(DINOv2ImageEncoder, self).__init__()
+
+        # Map model variant to feature dimension
+        variant_dims = {
+            'dinov2_vits14': 384,
+            'dinov2_vitb14': 768,
+            'dinov2_vitl14': 1024,
+            'dinov2_vitg14': 1536
+        }
+
+        if model_variant not in variant_dims:
+            raise ValueError(f"Unknown DINOv2 variant: {model_variant}")
+
+        self.feat_dim = variant_dims[model_variant]
+
+        # Load DINOv2 model
+        self.backbone = torch.hub.load('facebookresearch/dinov2', model_variant)
+
+        # Freeze backbone if requested
+        if freeze_backbone:
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
+    def forward(self, images, persistence_coords=None, pixel_coords=None):
+        """
+        Args:
+            images: (B, 3, H, W) tensor
+            persistence_coords: ignored (for interface compatibility)
+            pixel_coords: ignored (for interface compatibility)
+        Returns:
+            features: (B, feat_dim) raw DINOv2 features (384 for vits14)
+        """
+        # DINOv2 forward - return raw features
+        features = self.backbone(images)  # (B, feat_dim)
+
+        return features

graphformer/models/tree_encoder.py

@@ -0,0 +1,594 @@
+"""
+Tree Encoder Components for Neuron Morphology Analysis
+
+This module contains TreeLSTM and related components for encoding tree structures.
+"""
+
+import torch
+import torch.nn as nn
+import dgl
+
+
+class MultiNodeAggregation(nn.Module):
+    """Attention-based aggregation over all tree nodes instead of just root"""
+    def __init__(self, h_size, aggregation_type="attention"):
+        super(MultiNodeAggregation, self).__init__()
+        self.h_size = h_size
+        self.aggregation_type = aggregation_type
+
+        if aggregation_type == "attention":
+            # Learnable attention over nodes
+            self.attention = nn.Sequential(
+                nn.Linear(h_size, h_size),
+                nn.Tanh(),
+                nn.Linear(h_size, 1)
+            )
+        elif aggregation_type == "weighted":
+            # Learnable weighted sum
+            self.weight_net = nn.Sequential(
+                nn.Linear(h_size, h_size // 2),
+                nn.ReLU(),
+                nn.Linear(h_size // 2, 1),
+                nn.Sigmoid()
+            )
+
+    def forward(self, g, node_features, offsets):
+        """
+        Args:
+            g: DGL graph
+            node_features: (N, h_size) features for all nodes
+            offsets: indices of root nodes for each tree in batch
+
+        Returns:
+            aggregated: (batch_size, h_size) aggregated features
+        """
+        batch_size = len(offsets)
+        aggregated = []
+
+        # For each tree in batch
+        for i in range(batch_size):
+            # Get start and end indices for this tree's nodes
+            if i < batch_size - 1:
+                start_idx = offsets[i] if i == 0 else offsets[i-1]
+                end_idx = offsets[i+1]
+            else:
+                start_idx = offsets[i-1] if i > 0 else 0
+                end_idx = len(node_features)
+
+            # Extract this tree's node features
+            tree_feats = node_features[start_idx:end_idx]  # (num_nodes_i, h_size)
+
+            if self.aggregation_type == "attention":
+                # Attention weights
+                attn_scores = self.attention(tree_feats)  # (num_nodes_i, 1)
+                attn_weights = torch.softmax(attn_scores, dim=0)
+
+                # Weighted sum
+                agg = (tree_feats * attn_weights).sum(dim=0)  # (h_size,)
+
+            elif self.aggregation_type == "weighted":
+                # Simple learnable weighting
+                weights = self.weight_net(tree_feats)  # (num_nodes_i, 1)
+                agg = (tree_feats * weights).sum(dim=0)  # (h_size,)
+
+            elif self.aggregation_type == "mean":
+                agg = tree_feats.mean(dim=0)  # (h_size,)
+
+            elif self.aggregation_type == "max":
+                agg = tree_feats.max(dim=0)[0]  # (h_size,)
+
+            aggregated.append(agg)
+
+        return torch.stack(aggregated, dim=0)  # (batch_size, h_size)
+
+
+class TreeLSTMCell(nn.Module):
+    def __init__(self, x_size, h_size, mode="sum"):
+        super(TreeLSTMCell, self).__init__()
+        self.h_size, self.mode = h_size, mode
+        self.W_iou = nn.Linear(x_size, 3 * h_size, bias=False)
+        self.U_iou = nn.Linear(h_size, 3 * h_size, bias=False)
+        self.b_iou = nn.Parameter(torch.zeros(1, 3 * h_size))
+        self.U_f = nn.Linear(h_size, h_size)
+
+    def message_func(self, edges):
+        return {"h": edges.src["h"], "c": edges.src["c"]}
+
+    def reduce_func(self, nodes):
+        if self.mode == "sum":
+            h_cat = nodes.mailbox["h"].sum(dim=1)
+        elif self.mode == "max":
+            h_cat = nodes.mailbox["h"].max(dim=1)[0]
+        elif self.mode == "mean":
+            h_cat = nodes.mailbox["h"].mean(dim=1)
+        else:
+            raise NotImplementedError
+
+        f = torch.sigmoid(self.U_f(nodes.mailbox["h"]))
+        c = torch.sum(f * nodes.mailbox["c"], 1)
+        return {"iou": self.U_iou(h_cat), "c": c}
+
+    def apply_node_func(self, nodes):
+        iou = nodes.data["iou"] + self.b_iou
+        i, o, u = torch.chunk(iou, 3, 1)
+        i, o, u = torch.sigmoid(i), torch.sigmoid(o), torch.tanh(u)
+        c = i * u + nodes.data["c"]
+        h = o * torch.tanh(c)
+        return {"h": h, "c": c}
+
+
+class BidirectionalTreeLSTMCell(nn.Module):
+    """Bidirectional TreeLSTM Cell that processes tree in both bottom-up and top-down directions"""
+    def __init__(self, x_size, h_size, mode="sum"):
+        super(BidirectionalTreeLSTMCell, self).__init__()
+        self.h_size = h_size
+        self.mode = mode
+
+        # Bottom-up (child to parent) parameters
+        self.W_iou_bu = nn.Linear(x_size, 3 * h_size, bias=False)
+        self.U_iou_bu = nn.Linear(h_size, 3 * h_size, bias=False)
+        self.b_iou_bu = nn.Parameter(torch.zeros(1, 3 * h_size))
+        self.U_f_bu = nn.Linear(h_size, h_size)
+
+        # Top-down (parent to child) parameters
+        self.W_iou_td = nn.Linear(x_size, 3 * h_size, bias=False)
+        self.U_iou_td = nn.Linear(h_size, 3 * h_size, bias=False)
+        self.b_iou_td = nn.Parameter(torch.zeros(1, 3 * h_size))
+        self.U_f_td = nn.Linear(h_size, h_size)
+
+    def message_func_bu(self, edges):
+        """Bottom-up message function (from children to parent)"""
+        return {"h_bu": edges.src["h_bu"], "c_bu": edges.src["c_bu"]}
+
+    def reduce_func_bu(self, nodes):
+        """Bottom-up reduce function"""
+        if self.mode == "sum":
+            h_cat = nodes.mailbox["h_bu"].sum(dim=1)
+        elif self.mode == "max":
+            h_cat = nodes.mailbox["h_bu"].max(dim=1)[0]
+        elif self.mode == "mean":
+            h_cat = nodes.mailbox["h_bu"].mean(dim=1)
+        else:
+            raise NotImplementedError
+
+        f = torch.sigmoid(self.U_f_bu(nodes.mailbox["h_bu"]))
+        c = torch.sum(f * nodes.mailbox["c_bu"], 1)
+        return {"iou_bu": self.U_iou_bu(h_cat), "c_bu": c}
+
+    def apply_node_func_bu(self, nodes):
+        """Bottom-up apply node function"""
+        iou = nodes.data["iou_bu"] + self.b_iou_bu
+        i, o, u = torch.chunk(iou, 3, 1)
+        i, o, u = torch.sigmoid(i), torch.sigmoid(o), torch.tanh(u)
+        c = i * u + nodes.data["c_bu"]
+        h = o * torch.tanh(c)
+        return {"h_bu": h, "c_bu": c}
+
+    def message_func_td(self, edges):
+        """Top-down message function (from parent to children)"""
+        return {"h_td": edges.dst["h_td"], "c_td": edges.dst["c_td"]}
+
+    def reduce_func_td(self, nodes):
+        """Top-down reduce function"""
+        if nodes.mailbox["h_td"].shape[1] == 0:
+            # Root node has no parent
+            return {"iou_td": torch.zeros(nodes.batch_size(), 3 * self.h_size, device=nodes.mailbox["h_td"].device),
+                    "c_td": torch.zeros(nodes.batch_size(), self.h_size, device=nodes.mailbox["h_td"].device)}
+
+        # For non-root nodes, aggregate parent information
+        h_parent = nodes.mailbox["h_td"][:, 0, :]  # Take first (and only) parent
+        c_parent = nodes.mailbox["c_td"][:, 0, :]
+
+        return {"iou_td": self.U_iou_td(h_parent), "c_td": c_parent}
+
+    def apply_node_func_td(self, nodes):
+        """Top-down apply node function"""
+        iou = nodes.data["iou_td"] + self.b_iou_td
+        i, o, u = torch.chunk(iou, 3, 1)
+        i, o, u = torch.sigmoid(i), torch.sigmoid(o), torch.tanh(u)
+        c = i * u + nodes.data["c_td"]
+        h = o * torch.tanh(c)
+        return {"h_td": h, "c_td": c}
+
+
+class BidirectionalTreeLSTM(nn.Module):
+    """Bidirectional TreeLSTM that combines bottom-up and top-down processing"""
+    def __init__(self, x_size, h_size, num_classes, mode="sum", fc=True, bn=False,
+                 node_aggregation=None):
+        super(BidirectionalTreeLSTM, self).__init__()
+        self.x_size = x_size
+        self.h_size = h_size
+        self.node_aggregation = node_aggregation
+
+        # Input projection
+        if bn:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(x_size, h_size),
+                nn.BatchNorm1d(h_size),
+                nn.ReLU(),
+            )
+        else:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(x_size, h_size),
+                nn.ReLU(),
+            )
+
+        # Bidirectional TreeLSTM cell
+        self.cell = BidirectionalTreeLSTMCell(h_size, h_size, mode=mode)
+
+        # Node aggregation
+        if node_aggregation:
+            self.node_agg = MultiNodeAggregation(h_size * 2, aggregation_type=node_aggregation)
+
+        # Classification head
+        self.fc = fc
+        if fc:
+            self.linear = nn.Linear(h_size * 2, num_classes)
+
+    def forward(self, batch):
+        """Forward pass combining bottom-up and top-down TreeLSTM"""
+        g = batch.graph.to(torch.device("cuda"))
+        g = dgl.graph(g.edges())
+        n = g.number_of_nodes()
+
+        # Input projection
+        feats = self.mlp1(batch.feats.cuda())
+
+        # Initialize node states for both directions
+        g.ndata["iou_bu"] = self.cell.W_iou_bu(feats)
+        g.ndata["iou_td"] = self.cell.W_iou_td(feats)
+        g.ndata["h_bu"] = torch.zeros((n, self.h_size)).cuda()
+        g.ndata["c_bu"] = torch.zeros((n, self.h_size)).cuda()
+        g.ndata["h_td"] = torch.zeros((n, self.h_size)).cuda()
+        g.ndata["c_td"] = torch.zeros((n, self.h_size)).cuda()
+
+        # Bottom-up pass
+        dgl.prop_nodes_topo(
+            g,
+            message_func=self.cell.message_func_bu,
+            reduce_func=self.cell.reduce_func_bu,
+            apply_node_func=self.cell.apply_node_func_bu,
+        )
+
+        # Top-down pass (reverse topological order)
+        dgl.prop_nodes_topo(
+            g,
+            message_func=self.cell.message_func_td,
+            reduce_func=self.cell.reduce_func_td,
+            apply_node_func=self.cell.apply_node_func_td,
+            reverse=True,
+        )
+
+        # Combine bottom-up and top-down representations
+        h_combined = torch.cat([g.ndata.pop("c_bu"), g.ndata.pop("c_td")], dim=1)
+
+        # Aggregate node features
+        if self.node_aggregation:
+            h = self.node_agg(g, h_combined, batch.offset.long())
+        else:
+            h = h_combined[batch.offset.long()]
+
+        # Classification
+        if self.fc:
+            return self.linear(h)
+        return h
+
+
+class TreeLSTM(nn.Module):
+    def __init__(self, x_size, h_size, num_classes, mode="sum", fc=True, bn=False,
+                 node_aggregation=None):
+        super(TreeLSTM, self).__init__()
+        self.x_size = x_size
+        self.h_size = h_size
+        self.node_aggregation = node_aggregation
+
+        if bn:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(x_size, h_size),
+                nn.BatchNorm1d(h_size),
+                nn.ReLU(),
+            )
+        else:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(x_size, h_size),
+                nn.ReLU(),
+            )
+
+        self.cell = TreeLSTMCell(h_size, h_size, mode=mode)
+
+        if node_aggregation:
+            self.node_agg = MultiNodeAggregation(h_size, aggregation_type=node_aggregation)
+
+        self.fc = fc
+        if fc:
+            self.linear = nn.Linear(h_size, num_classes)
+
+    def forward(self, batch):
+        g = batch.graph.to(torch.device("cuda"))
+        g = dgl.graph(g.edges())
+        n = g.number_of_nodes()
+
+        feats = self.mlp1(batch.feats.cuda())
+        g.ndata["iou"] = self.cell.W_iou(feats)
+        g.ndata["h"] = torch.zeros((n, self.h_size)).cuda()
+        g.ndata["c"] = torch.zeros((n, self.h_size)).cuda()
+
+        dgl.prop_nodes_topo(
+            g,
+            message_func=self.cell.message_func,
+            reduce_func=self.cell.reduce_func,
+            apply_node_func=self.cell.apply_node_func,
+        )
+
+        h = g.ndata.pop("c")
+
+        if self.node_aggregation:
+            h = self.node_agg(g, h, batch.offset.long())
+        else:
+            h = h[batch.offset.long()]
+
+        if self.fc:
+            return self.linear(h)
+        return h
+
+
+class TreeLSTM_wo_MLP(nn.Module):
+    """TreeLSTM without initial MLP projection"""
+    def __init__(self, x_size, h_size, num_classes, mode="sum", fc=True):
+        super(TreeLSTM_wo_MLP, self).__init__()
+        self.x_size = x_size
+        self.h_size = h_size
+        self.cell = TreeLSTMCell(x_size, h_size, mode=mode)
+        self.fc = fc
+        if fc:
+            self.linear = nn.Linear(h_size, num_classes)
+
+    def forward(self, batch):
+        g = batch.graph.to(torch.device("cuda"))
+        g = dgl.graph(g.edges())
+        n = g.number_of_nodes()
+
+        feats = batch.feats.cuda()
+        g.ndata["iou"] = self.cell.W_iou(feats)
+        g.ndata["h"] = torch.zeros((n, self.h_size)).cuda()
+        g.ndata["c"] = torch.zeros((n, self.h_size)).cuda()
+
+        dgl.prop_nodes_topo(
+            g,
+            message_func=self.cell.message_func,
+            reduce_func=self.cell.reduce_func,
+            apply_node_func=self.cell.apply_node_func,
+        )
+
+        h = g.ndata.pop("c")[batch.offset.long()]
+
+        if self.fc:
+            return self.linear(h)
+        return h
+
+
+class TreeLSTMCellv2(nn.Module):
+    """TreeLSTM Cell with alternative architecture"""
+    def __init__(self, x_size, h_size, mode="sum"):
+        super(TreeLSTMCellv2, self).__init__()
+        self.h_size = h_size
+        self.mode = mode
+        self.W_iou = nn.Linear(x_size, 3 * h_size, bias=False)
+        self.U_iou = nn.Linear(2 * h_size, 3 * h_size, bias=False)
+        self.b_iou = nn.Parameter(torch.zeros(1, 3 * h_size))
+        self.U_f = nn.Linear(2 * h_size, h_size)
+
+    def message_func(self, edges):
+        return {"h": edges.src["h"], "c": edges.src["c"]}
+
+    def reduce_func(self, nodes):
+        if self.mode == "sum":
+            h_cat = nodes.mailbox["h"].sum(dim=1)
+        elif self.mode == "max":
+            h_cat = nodes.mailbox["h"].max(dim=1)[0]
+        elif self.mode == "mean":
+            h_cat = nodes.mailbox["h"].mean(dim=1)
+
+        h_max = nodes.mailbox["h"].max(dim=1)[0]
+        h_combined = torch.cat([h_cat, h_max], dim=1)
+
+        f = torch.sigmoid(self.U_f(h_combined.unsqueeze(1).expand(-1, nodes.mailbox["h"].shape[1], -1)))
+        c = torch.sum(f * nodes.mailbox["c"], 1)
+        return {"iou": self.U_iou(h_combined), "c": c}
+
+    def apply_node_func(self, nodes):
+        iou = nodes.data["iou"] + self.b_iou
+        i, o, u = torch.chunk(iou, 3, 1)
+        i, o, u = torch.sigmoid(i), torch.sigmoid(o), torch.tanh(u)
+        c = i * u + nodes.data["c"]
+        h = o * torch.tanh(c)
+        return {"h": h, "c": c}
+
+
+class TreeLSTMDoubleCell(nn.Module):
+    """TreeLSTM Cell with double hidden state (original implementation)"""
+    def __init__(self, x_size, h_size, mode="sum"):
+        super(TreeLSTMDoubleCell, self).__init__()
+        self.W1_iouf = nn.Linear(x_size, 4 * h_size)
+        self.U1_iouf = nn.Linear(h_size, 4 * h_size)
+        self.W2_iouf = nn.Linear(h_size, 4 * h_size)
+        self.U2_iouf = nn.Linear(h_size, 4 * h_size)
+        self.mode = mode
+        self.init_state = True
+        self.h_size = h_size
+
+    def message_func(self, edges):
+        return {
+            "h1": edges.src["h1"],
+            "c1": edges.src["c1"],
+            "h2": edges.src["h2"],
+            "c2": edges.src["c2"],
+        }
+
+    def reduce_func(self, nodes):
+        h1, c1 = nodes.mailbox["h1"], nodes.mailbox["c1"]
+        h2, c2 = nodes.mailbox["h2"], nodes.mailbox["c2"]
+        if self.mode == "sum":
+            h1, c1, h2, c2 = h1.sum(-2), c1.sum(-2), h2.sum(-2), c2.sum(-2)
+        elif self.mode == "mean":
+            h1, c1, h2, c2 = h1.mean(-2), c1.mean(-2), h2.mean(-2), c2.mean(-2)
+        else:
+            raise ValueError("must in [sum, mean]")
+        x_iouf = nodes.data["iouf"]
+        xi, xo, xu, xf = torch.chunk(x_iouf, 4, 1)
+        h_iouf1 = self.U1_iouf(h1)
+        hi1, ho1, hu1, hf1 = torch.chunk(h_iouf1, 4, 1)
+        i = torch.sigmoid(xi + hi1)
+        f = torch.sigmoid(xf + hf1)
+        o = torch.sigmoid(xo + ho1)
+        u = torch.tanh(xu + hu1)
+        c1 = i * u + f * c1
+        h1 = o * torch.tanh(c1)
+
+        x_iouf2 = self.W2_iouf(c1)
+        xi, xo, xu, xf = torch.chunk(x_iouf2, 4, 1)
+        h_iouf2 = self.U2_iouf(h2)
+        hi2, ho2, hu2, hf2 = torch.chunk(h_iouf2, 4, 1)
+        i = torch.sigmoid(xi + hi2)
+        f = torch.sigmoid(xf + hf2)
+        o = torch.sigmoid(xo + ho2)
+        u = torch.tanh(xu + hu2)
+        c2 = i * u + f * c2
+        h2 = o * torch.tanh(c2)
+        return {"h1": h1, "c1": c1, "h2": h2, "c2": c2}
+
+    def apply_node_func(self, nodes):
+        if self.init_state:
+            iouf = nodes.data["iouf"]
+            i, o, u, f = torch.chunk(iouf, 4, 1)
+            i, o, u, f = torch.sigmoid(i), torch.sigmoid(o), torch.tanh(u), torch.sigmoid(f)
+            c1 = i * u + f * nodes.data["c1"]
+            h1 = o * torch.tanh(c1)
+
+            iouf2 = self.W2_iouf(c1)
+            i, o, u, f = torch.chunk(iouf2, 4, 1)
+            i, o, u, f = torch.sigmoid(i), torch.sigmoid(o), torch.tanh(u), torch.sigmoid(f)
+            c2 = i * u + f * nodes.data["c2"]
+            h2 = o * torch.tanh(c2)
+            self.init_state = False
+            return {"h1": h1, "c1": c1, "h2": h2, "c2": c2}
+        else:
+            return {
+                "h1": nodes.data["h1"],
+                "c1": nodes.data["c1"],
+                "h2": nodes.data["h2"],
+                "c2": nodes.data["c2"],
+            }
+
+
+class TreeLSTMDouble(nn.Module):
+    """TreeLSTM with double hidden state aggregation"""
+    def __init__(self, x_size, h_size, num_classes, mode="sum", fc=True, bn=False, node_aggregation=None):
+        super(TreeLSTMDouble, self).__init__()
+        self.x_size, self.h_size = x_size, h_size
+        self.node_aggregation = node_aggregation
+
+        if bn:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(x_size, h_size),
+                nn.BatchNorm1d(h_size),
+                nn.ReLU(),
+                nn.Linear(h_size, 2 * h_size),
+                nn.BatchNorm1d(2 * h_size),
+                nn.ReLU(),
+                nn.Linear(2 * h_size, h_size),
+            )
+        else:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(x_size, h_size),
+                nn.ReLU(),
+                nn.Linear(h_size, 2 * h_size),
+                nn.ReLU(),
+                nn.Linear(2 * h_size, h_size),
+            )
+
+        self.cell = TreeLSTMDoubleCell(h_size, h_size, mode=mode)
+
+        if node_aggregation:
+            self.node_agg = MultiNodeAggregation(h_size, aggregation_type=node_aggregation)
+
+        self.fc = fc
+        if fc:
+            self.linear = nn.Linear(h_size, num_classes)
+
+    def forward_backbone(self, batch):
+        g = batch.graph.to(torch.device("cuda"))
+        # to heterogenous graph
+        g = dgl.graph(g.edges())
+        n = g.number_of_nodes()
+        # feed embedding
+        feats = self.mlp1(batch.feats.cuda())
+        g.ndata["iouf"] = self.cell.W1_iouf(feats)
+        g.ndata["h1"] = torch.zeros((n, self.h_size)).cuda()
+        g.ndata["c1"] = torch.zeros((n, self.h_size)).cuda()
+        g.ndata["h2"] = torch.zeros((n, self.h_size)).cuda()
+        g.ndata["c2"] = torch.zeros((n, self.h_size)).cuda()
+        # propagate
+        dgl.prop_nodes_topo(
+            g,
+            message_func=self.cell.message_func,
+            reduce_func=self.cell.reduce_func,
+            apply_node_func=self.cell.apply_node_func,
+        )
+        logits = g.ndata.pop("c2")[batch.offset.long()]
+        return logits
+
+    def forward(self, batch):
+        logits = self.forward_backbone(batch)
+        if self.fc:
+            logits = self.linear(logits)
+            return logits
+        else:
+            return logits
+
+
+class TreeLSTMv2(nn.Module):
+    """TreeLSTM variant with improved architecture"""
+    def __init__(self, x_size, h_size, num_classes, mode="sum", fc=True, bn=False):
+        super(TreeLSTMv2, self).__init__()
+        self.x_size, self.h_size = x_size, h_size
+
+        if bn:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(x_size, h_size),
+                nn.BatchNorm1d(h_size),
+                nn.ReLU(),
+            )
+        else:
+            self.mlp1 = nn.Sequential(
+                nn.Linear(x_size, h_size),
+                nn.ReLU(),
+            )
+
+        self.cell = TreeLSTMCellv2(h_size, h_size, mode=mode)
+        self.fc = fc
+        if fc:
+            self.linear = nn.Linear(h_size, num_classes)
+
+    def forward(self, batch):
+        g = batch.graph.to(torch.device("cuda"))
+        g = dgl.graph(g.edges())
+        n = g.number_of_nodes()
+
+        feats = self.mlp1(batch.feats.cuda())
+        g.ndata["iou"] = self.cell.W_iou(feats)
+        g.ndata["h"] = torch.zeros((n, self.h_size)).cuda()
+        g.ndata["c"] = torch.zeros((n, self.h_size)).cuda()
+
+        dgl.prop_nodes_topo(
+            g,
+            message_func=self.cell.message_func,
+            reduce_func=self.cell.reduce_func,
+            apply_node_func=self.cell.apply_node_func,
+        )
+
+        h = g.ndata.pop("c")[batch.offset.long()]
+
+        if self.fc:
+            return self.linear(h)
+        return h

graphformer/utils/__init__.py

@@ -0,0 +1,7 @@
+"""Utility functions for GraPHFormer."""
+
+from .training import set_seed, save_checkpoint, adjust_learning_rate, get_root_logger
+
+__all__ = [
+    "set_seed", "save_checkpoint", "adjust_learning_rate", "get_root_logger",
+]

graphformer/utils/training.py

@@ -0,0 +1,102 @@
+import torch
+import shutil
+import math
+import logging
+import numpy as np
+import random
+
+
+def set_seed(seed):
+    """
+    Setting of Global Seed
+
+    """
+    torch.backends.cudnn.enabled = True
+    torch.backends.cudnn.deterministic = True  # consistent results on the cpu and gpu
+    torch.backends.cudnn.benchmark = False
+
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)  # cpu
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # gpu
+
+
+def save_checkpoint(state, is_best, filename="checkpoint.pth"):
+    torch.save(state, filename)
+    if is_best:
+        pth = "/".join(filename.split("/")[:-1])
+        shutil.copyfile(filename, f"{pth}/model_best.pth")
+
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """Decay the learning rate based on schedule"""
+    lr = args.lr
+    if args.cos:  # cosine lr schedule
+        lr *= 0.5 * (1.0 + math.cos(math.pi * epoch / args.epochs))
+    else:  # stepwise lr schedule
+        for milestone in args.schedule:
+            lr *= 0.1 if epoch >= milestone else 1.0
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+
+def get_root_logger(log_file=None, log_level=logging.INFO):
+    """Get the root logger.
+
+    The logger will be initialized if it has not been initialized. By default a
+    StreamHandler will be added. If `log_file` is specified, a FileHandler will
+    also be added. The name of the root logger is the top-level package name,
+    e.g., "openselfsup".
+
+    Args:
+        log_file (str | None): The log filename. If specified, a FileHandler
+            will be added to the root logger.
+        log_level (int): The root logger level. Note that only the process of
+            rank 0 is affected, while other processes will set the level to
+            "Error" and be silent most of the time.
+
+    Returns:
+        logging.Logger: The root logger.
+    """
+    logger = logging.getLogger(__name__.split(".")[0])  # i.e., openselfsup
+    # if the logger has been initialized, just return it
+    if logger.hasHandlers():
+        return logger
+
+    format_str = "%(asctime)s - %(message)s"
+    logging.basicConfig(format=format_str, level=log_level)
+    if log_file is not None:
+        file_handler = logging.FileHandler(log_file, "w")
+        file_handler.setFormatter(logging.Formatter(format_str))
+        file_handler.setLevel(log_level)
+        logger.addHandler(file_handler)
+
+    return logger
+
+
+def print_log(msg, logger=None, level=logging.INFO):
+    """Print a log message.
+
+    Args:
+        msg (str): The message to be logged.
+        logger (logging.Logger | str | None): The logger to be used. Some
+            special loggers are:
+            - "root": the root logger obtained with `get_root_logger()`.
+            - "silent": no message will be printed.
+            - None: The `print()` method will be used to print log messages.
+        level (int): Logging level. Only available when `logger` is a Logger
+            object or "root".
+    """
+    if logger is None:
+        print(msg)
+    elif logger == "root":
+        _logger = get_root_logger()
+        _logger.log(level, msg)
+    elif isinstance(logger, logging.Logger):
+        logger.log(level, msg)
+    elif logger != "silent":
+        raise TypeError(
+            'logger should be either a logging.Logger object, "root", '
+            '"silent" or None, but got {}'.format(logger)
+        )

scripts/finetune_example.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+# Example fine-tuning script for GraPHFormer
+
+python finetune.py \
+    --exp_name finetune_bil \
+    --dataset bil_6_classes \
+    --pretrained_checkpoint work_dir/graphformer_resnet18/model_best.pth \
+    --mode multimodal \
+    --fusion_mode concat \
+    --tree_model double \
+    --image_encoder dinov2_vits14 \
+    --image_size 252 \
+    --embed_dim 128 \
+    --h_size 256 \
+    --batch_size 64 \
+    --epochs 50 \
+    --lr 1e-4 \
+    --wd 0.01 \
+    --dropout 0.5 \
+    --label_smoothing 0.1 \
+    --linear_probe_epochs 10 \
+    --early_stopping_patience 10 \
+    --val_freq 1 \
+    --eval_mode accuracy

scripts/prepare_data.py

@@ -0,0 +1,236 @@
+"""
+Please refer to data/README.md for data download and preparation.
+Then run this script to preprocess data.
+"""
+
+import os, sys
+from data_io.process_raw_script import (
+    convert_janelia_json,
+    normalize_root_and_check,
+    filter_axon_and_check,
+    summarize_branch,
+    split_sample_10fold_cv_and_merge,
+)
+import pandas as pd
+
+data_path = "./data"
+
+# convert JML json file
+convert_janelia_json(
+    os.path.join(data_path, "raw/janelia_mouselight/json30/*.json"),
+    os.path.join(data_path, "raw/janelia_mouselight/swc"),
+    os.path.join(data_path, "info/JML_info_swc.csv"),
+)
+
+# preprocess data
+for source in ["janelia_mouselight", "allen_cell_type", "bil"]:
+    print(f"Processing:{source}")
+    # normalize neuron's center, orientation and size
+    print(f"Normalize neuron")
+    folder_in = f"{data_path}/raw/{source}/swc/"
+    folder_out = f"{data_path}/raw/{source}/swc_soma0/"
+    if source == "bil":
+        normalize_root_and_check(folder_in + "*reg.swc", folder_out)
+        # some BIL reconstructions are not correctly scaled. this will fix them
+        normalize_root_and_check(
+            folder_in + "*__reg.swc", folder_out, scale=[0.114, 0.114, 0.28]
+        )
+    else:
+        normalize_root_and_check(folder_in + "*.swc", folder_out)
+
+    # remove axon file
+    print(f"Remove axons")
+    folder_in = f"{data_path}/raw/{source}/swc_soma0/*.swc"
+    folder_out = f"{data_path}/dendrite/{source}/swc_soma0/"
+    filter_axon_and_check(folder_in, folder_out)
+
+    print(f"Calculate features")
+    folder_in = f"{data_path}/dendrite/{source}/swc_soma0/*.swc"
+    folder_out = f"{data_path}/dendrite/{source}/eswc_soma0/"
+    summarize_branch(folder_in, folder_out)
+
+
+# split data into 10 folds
+split_sample_10fold_cv_and_merge(data_path)
+
+# hack folder creation
+folder_names = ["allen_cell_type", "bil", "janelia_mouselight"]
+all_wo_others = {
+    "VPM": 0,
+    "Isocortex_layer23": 1,
+    "Isocortex_layer4": 2,
+    "PRE": 3,
+    "SUB": 4,
+    "CP": 5,
+    "VPL": 6,
+    "Isocortex_layer6": 7,
+    "MG": 8,
+    "Isocortex_layer5": 9,
+}
+for i, split_csv in enumerate(
+    [
+        f"{data_path}/info/ACT_info_swc_10folds.csv",
+        f"{data_path}/info/BIL_info_swc_10folds.csv",
+        f"{data_path}/info/JML_info_swc_10folds.csv",
+    ]
+):
+    csv = pd.read_csv(split_csv)
+    folder_name = folder_names[i]
+    for split in range(10):
+        for fname in csv[csv["model__fold"] == split]["swc__fname"]:
+            # get acronym from "structure_merge__acronym"
+            acronym = csv[csv["swc__fname"] == fname][
+                "structure_merge__acronym"
+            ].values[0]
+            if acronym not in all_wo_others:
+                if acronym == "Isocortex_layer2/3":
+                    acronym = "Isocortex_layer23"
+                else:
+                    continue
+            os.makedirs(
+                f"{data_path}/dendrite/all_eswc_soma0_ssl/{acronym}/{folder_name}-{split}/",
+                exist_ok=True,
+            )
+            source_path = os.path.abspath(
+                f"{data_path}/dendrite/{folder_name}/eswc_soma0/{fname}"
+            )
+            target_path = f"{data_path}/dendrite/all_eswc_soma0_ssl/{acronym}/{folder_name}-{split}/{fname}"
+
+            # Check if the source path is a valid file
+            if os.path.isfile(source_path):
+                os.symlink(source_path, target_path)
+            else:
+                print(f"{source_path} is not a valid file!")
+
+
+# # hack folder creation
+# folder_names = ["allen_cell_type"]
+# all_wo_others = {
+#     "Isocortex_layer23": 0,
+#     "Isocortex_layer4": 1,
+#     "Isocortex_layer5": 2,
+#     "Isocortex_layer6": 3,
+# }
+# for i, split_csv in enumerate(
+#     [
+#         f"{data_path}/info/ACT_info_swc_10folds.csv",
+#         # f"{data_path}/info/BIL_info_swc_10folds.csv",
+#         # f"{data_path}/info/JML_info_swc_10folds.csv",
+#     ]
+# ):
+#     csv = pd.read_csv(split_csv)
+#     folder_name = folder_names[i]
+#     for split in range(10):
+#         for fname in csv[csv["model__fold"] == split]["swc__fname"]:
+#             # get acronym from "structure_merge__acronym"
+#             acronym = csv[csv["swc__fname"] == fname][
+#                 "structure_merge__acronym"
+#             ].values[0]
+#             if acronym not in all_wo_others:
+#                 if acronym == "Isocortex_layer2/3":
+#                     acronym = "Isocortex_layer23"
+#                 else:
+#                     continue
+#             os.makedirs(
+#                 f"{data_path}/dendrite/ACT/{acronym}/{folder_name}-{split}/",
+#                 exist_ok=True,
+#             )
+#             source_path = os.path.abspath(
+#                 f"{data_path}/dendrite/{folder_name}/eswc_soma0/{fname}"
+#             )
+#             target_path = f"{data_path}/dendrite/ACT/{acronym}/{folder_name}-{split}/{fname}"
+
+#             # Check if the source path is a valid file
+#             if os.path.isfile(source_path):
+#                 os.symlink(source_path, target_path)
+#             else:
+#                 print(f"{source_path} is not a valid file!")
+
+
+# # hack folder creation
+# folder_names = ["janelia_mouselight"]
+# all_wo_others = {
+#     "Isocortex_layer23": 0,
+#     "Isocortex_layer5": 1,
+#     "Isocortex_layer6": 2,
+#     "VPM": 3,
+# }
+# for i, split_csv in enumerate(
+#     [
+#         # f"{data_path}/info/ACT_info_swc_10folds.csv",
+#         # f"{data_path}/info/BIL_info_swc_10folds.csv",
+#         f"{data_path}/info/JML_info_swc_10folds.csv",
+#     ]
+# ):
+#     csv = pd.read_csv(split_csv)
+#     folder_name = folder_names[i]
+#     for split in range(10):
+#         for fname in csv[csv["model__fold"] == split]["swc__fname"]:
+#             # get acronym from "structure_merge__acronym"
+#             acronym = csv[csv["swc__fname"] == fname][
+#                 "structure_merge__acronym"
+#             ].values[0]
+#             if acronym not in all_wo_others:
+#                 if acronym == "Isocortex_layer2/3":
+#                     acronym = "Isocortex_layer23"
+#                 else:
+#                     continue
+#             os.makedirs(
+#                 f"{data_path}/dendrite/JML/{acronym}/{folder_name}-{split}/",
+#                 exist_ok=True,
+#             )
+#             source_path = os.path.abspath(
+#                 f"{data_path}/dendrite/{folder_name}/eswc_soma0/{fname}"
+#             )
+#             target_path = f"{data_path}/dendrite/JML/{acronym}/{folder_name}-{split}/{fname}"
+
+#             # Check if the source path is a valid file
+#             if os.path.isfile(source_path):
+#                 os.symlink(source_path, target_path)
+#             else:
+#                 print(f"{source_path} is not a valid file!")
+
+# # hack folder creation
+# folder_names = ["bil"]
+# all_wo_others = {
+#     "CP": 0,
+#     "Isocortex_layer23": 1,
+#     "Isocortex_layer4": 2,
+#     "Isocortex_layer5": 3,
+#     "Isocortex_layer6": 4,
+#     "VPM": 5,
+# }
+# for i, split_csv in enumerate(
+#     [
+#         # f"{data_path}/info/ACT_info_swc_10folds.csv",
+#         f"{data_path}/info/BIL_info_swc_10folds.csv",
+#         # f"{data_path}/info/JML_info_swc_10folds.csv",
+#     ]
+# ):
+#     csv = pd.read_csv(split_csv)
+#     folder_name = folder_names[i]
+#     for split in range(10):
+#         for fname in csv[csv["model__fold"] == split]["swc__fname"]:
+#             # get acronym from "structure_merge__acronym"
+#             acronym = csv[csv["swc__fname"] == fname][
+#                 "structure_merge__acronym"
+#             ].values[0]
+#             if acronym not in all_wo_others:
+#                 if acronym == "Isocortex_layer2/3":
+#                     acronym = "Isocortex_layer23"
+#                 else:
+#                     continue
+#             os.makedirs(
+#                 f"{data_path}/dendrite/BIL/{acronym}/{folder_name}-{split}/",
+#                 exist_ok=True,
+#             )
+#             source_path = os.path.abspath(
+#                 f"{data_path}/dendrite/{folder_name}/eswc_soma0/{fname}"
+#             )
+#             target_path = f"{data_path}/dendrite/BIL/{acronym}/{folder_name}-{split}/{fname}"
+
+#             # Check if the source path is a valid file
+#             if os.path.isfile(source_path):
+#                 os.symlink(source_path, target_path)
+#             else:
+#                 print(f"{source_path} is not a valid file!")

scripts/train_example.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+# Example training script for GraPHFormer
+
+python train.py \
+    --exp_name graphformer_resnet18 \
+    --dataset all_wo_others \
+    --tree_model double \
+    --image_encoder dinov2_vits14 \
+    --image_size 252 \
+    --embed_dim 128 \
+    --h_size 256 \
+    --batch_size 128 \
+    --epochs 100 \
+    --lr 3e-4 \
+    --wd 0.1 \
+    --temperature 0.07 \
+    --loss_type clip \
+    --warmup_epochs 5 \
+    --save_freq 10 \
+    --val_freq 5 \
+    --aug_rotate \
+    --aug_jitter_coords \
+    --use_persistence_aug \
+    --use_knn_eval \
+    --knn_k 20 \
+    --eval_jm \
+    --eval_act

setup.py

@@ -0,0 +1,30 @@
+"""Setup script for GraPHFormer."""
+
+from setuptools import setup, find_packages
+
+setup(
+    name="graphformer",
+    version="1.0.0",
+    description="GraPHFormer: Graph-Persistence Hybrid Transformer for Neuron Morphology",
+    author="",
+    packages=find_packages(),
+    python_requires=">=3.8",
+    install_requires=[
+        "torch>=1.10.0",
+        "torchvision>=0.11.0",
+        "dgl>=0.8.0",
+        "numpy>=1.20.0",
+        "scikit-learn>=1.0.0",
+        "nltk>=3.6.0",
+        "tqdm>=4.60.0",
+        "networkx>=2.6.0",
+        "Pillow>=8.0.0",
+    ],
+    extras_require={
+        "dev": [
+            "pytest>=6.0.0",
+            "black>=21.0",
+            "flake8>=3.9.0",
+        ],
+    },
+)

train.py

@@ -0,0 +1,786 @@
+"""
+GraPHFormer Training Script
+
+CLIP-style contrastive learning for neuron morphology representation.
+Aligns tree structure representations with persistence images.
+
+Usage:
+    python train.py --exp_name my_experiment --dataset all_wo_others
+"""
+
+import argparse
+import datetime
+import time
+import os
+import json
+import copy
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.utils.data import DataLoader
+from sklearn.neighbors import KNeighborsClassifier
+
+from graphformer.models import CLIPModel
+from graphformer.augmentations import (
+    Compose,
+    RandomScaleCoords, RandomRotate, RandomJitter, RandomShift,
+    RandomFlip, RandomMaskFeats, RandomJitterLength, RandomElasticate,
+    RandomDropSubTrees, RandomSkipParentNode, RandomSwapSiblingSubTrees,
+    CombinedPersistenceAugmentation,
+)
+from graphformer.utils import save_checkpoint, get_root_logger, set_seed
+from graphformer.data import NeuronTreeDataset, get_collate_fn, LABEL_DICT
+
+
+def knn_predict(feature, feature_bank, feature_labels, classes, knn_k, knn_t):
+    """KNN classification"""
+    feature = F.normalize(feature, dim=-1)
+    feature_bank = F.normalize(feature_bank, dim=-1)
+
+    sim_matrix = torch.mm(feature, feature_bank.t())
+    sim_weight, sim_indices = sim_matrix.topk(k=knn_k, dim=-1)
+
+    sim_labels = torch.gather(
+        feature_labels.expand(feature.size(0), -1), dim=-1, index=sim_indices
+    )
+
+    sim_weight = (sim_weight / knn_t).exp()
+
+    one_hot_label = torch.zeros(
+        feature.size(0) * knn_k, classes, device=sim_labels.device
+    )
+    one_hot_label = one_hot_label.scatter(
+        dim=-1, index=sim_labels.view(-1, 1), value=1.0
+    )
+
+    pred_scores = torch.sum(
+        one_hot_label.view(feature.size(0), -1, classes) * sim_weight.unsqueeze(dim=-1),
+        dim=1,
+    )
+    pred_labels = pred_scores.argmax(dim=-1)
+
+    return pred_labels
+
+
+def get_features_from_encoder(model, data_loader, device, fusion='concat'):
+    """Extract features from encoder"""
+    model.eval()
+    features = []
+    labels = []
+
+    with torch.no_grad():
+        for batch in data_loader:
+            batch = batch.to(device)
+            tree_embed = model.encode_tree(batch)
+            images = batch.images.cuda() if not batch.images.is_cuda else batch.images
+            image_embed = model.encode_image(images)
+
+            tree_embed = F.normalize(tree_embed, dim=-1)
+            image_embed = F.normalize(image_embed, dim=-1)
+
+            if fusion == 'concat':
+                combined = torch.cat([tree_embed, image_embed], dim=1)
+            elif fusion == 'add':
+                combined = tree_embed + image_embed
+            elif fusion == 'tree_only':
+                combined = tree_embed
+            elif fusion == 'image_only':
+                combined = image_embed
+            else:
+                combined = torch.cat([tree_embed, image_embed], dim=1)
+
+            combined = F.normalize(combined, dim=-1)
+
+            features.append(combined)
+            labels.append(batch.label.to(device))
+
+    features = torch.cat(features, dim=0)
+    labels = torch.cat(labels, dim=0)
+
+    return features, labels
+
+
+def evaluate_sklearn_knn(model, train_loader, test_loader, device, knn_k=20, fusion='concat'):
+    """Sklearn KNN evaluation"""
+    tmp_model = copy.deepcopy(model).eval()
+
+    x_train, y_train = get_features_from_encoder(tmp_model, train_loader, device, fusion)
+    x_test, y_test = get_features_from_encoder(tmp_model, test_loader, device, fusion)
+
+    neigh = KNeighborsClassifier(n_neighbors=knn_k)
+    neigh.fit(x_train.cpu().numpy(), y_train.cpu().numpy())
+
+    score = neigh.score(x_test.cpu().numpy(), y_test.cpu().numpy())
+
+    del tmp_model
+    return score
+
+
+def evaluate_knn(model, memory_loader, test_loader, device, num_classes, knn_k=20, knn_t=0.5, fusion='concat'):
+    """KNN evaluation during training"""
+    model.eval()
+
+    # Build memory bank from training set
+    feature_bank = []
+    with torch.no_grad():
+        for batch in memory_loader:
+            batch = batch.to(device)
+            tree_embed = model.encode_tree(batch)
+            images = batch.images.cuda() if not batch.images.is_cuda else batch.images
+            image_embed = model.encode_image(images)
+
+            tree_embed = F.normalize(tree_embed, dim=-1)
+            image_embed = F.normalize(image_embed, dim=-1)
+
+            if fusion == 'concat':
+                combined = torch.cat([tree_embed, image_embed], dim=1)
+            elif fusion == 'add':
+                combined = tree_embed + image_embed
+            elif fusion == 'tree_only':
+                combined = tree_embed
+            elif fusion == 'image_only':
+                combined = image_embed
+            else:
+                combined = torch.cat([tree_embed, image_embed], dim=1)
+
+            combined = F.normalize(combined, dim=-1)
+            feature_bank.append(combined)
+
+    feature_bank = torch.cat(feature_bank, dim=0).t().contiguous()
+    feature_labels = torch.tensor(memory_loader.dataset.targets, device=device)
+
+    # Test
+    total_top1, total_num = 0.0, 0
+    with torch.no_grad():
+        for batch in test_loader:
+            batch = batch.to(device)
+            tree_embed = model.encode_tree(batch)
+            images = batch.images.cuda() if not batch.images.is_cuda else batch.images
+            image_embed = model.encode_image(images)
+
+            tree_embed = F.normalize(tree_embed, dim=-1)
+            image_embed = F.normalize(image_embed, dim=-1)
+
+            if fusion == 'concat':
+                combined = torch.cat([tree_embed, image_embed], dim=1)
+            elif fusion == 'add':
+                combined = tree_embed + image_embed
+            elif fusion == 'tree_only':
+                combined = tree_embed
+            elif fusion == 'image_only':
+                combined = image_embed
+            else:
+                combined = torch.cat([tree_embed, image_embed], dim=1)
+
+            combined = F.normalize(combined, dim=-1)
+
+            pred_labels = knn_predict(
+                combined, feature_bank.t(), feature_labels,
+                num_classes, knn_k, knn_t
+            )
+
+            total_num += combined.size(0)
+            total_top1 += (pred_labels == batch.label.to(device)).float().sum().item()
+
+    accuracy = total_top1 / total_num * 100
+    model.train()
+    return accuracy
+
+
+def create_eval_dataset(phase, dataset_name, args):
+    """Create evaluation dataset"""
+    return NeuronTreeDataset(
+        phase=phase,
+        dataset=dataset_name,
+        label_dict=LABEL_DICT[dataset_name],
+        input_features=args.input_features,
+        use_images=True,
+        image_size=args.image_size,
+        cache_images=True,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("GraPHFormer Training")
+
+    # Basic
+    parser.add_argument("--work_dir", type=str, default="./work_dir")
+    parser.add_argument("--exp_name", type=str, required=True)
+    parser.add_argument("--dataset", type=str, default="all_wo_others")
+    parser.add_argument("--data_dir", type=str, default="data/raw/bil")
+    parser.add_argument("--seed", type=int, default=42)
+
+    # Tree Model
+    parser.add_argument("--tree_model", type=str, default="double",
+                        choices=["ori", "v2", "double"])
+    parser.add_argument("--child_mode", type=str, default="sum")
+    parser.add_argument("--input_features", nargs="+", type=int,
+                        default=[2, 3, 4, 12, 13])
+    parser.add_argument("--h_size", type=int, default=256)
+    parser.add_argument("--bn", action="store_true", default=False)
+
+    # Image Model
+    parser.add_argument("--image_encoder", type=str, default="resnet18",
+                        choices=["resnet18", "resnet50", "resnet101", "simplecnn",
+                                "smallvit", "persistencevit", "dinov2_vits14"])
+    parser.add_argument("--image_size", type=int, default=256)
+    parser.add_argument("--freeze_image_backbone", action="store_true", default=False)
+
+    # CLIP settings
+    parser.add_argument("--embed_dim", type=int, default=128)
+    parser.add_argument("--single_linear_proj", action="store_true", default=False)
+    parser.add_argument("--temperature", type=float, default=0.07)
+    parser.add_argument("--loss_type", type=str, default="clip",
+                        choices=["clip", "infonce", "ntxent", "triplet"])
+
+    # Triplet loss parameters
+    parser.add_argument("--triplet_margin", type=float, default=1.0)
+    parser.add_argument("--triplet_mining", type=str, default="batch_hard")
+    parser.add_argument("--triplet_distance", type=str, default="euclidean")
+
+    # Training
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--epochs", default=100, type=int)
+    parser.add_argument("--lr", default=3e-4, type=float)
+    parser.add_argument("--wd", default=0.1, type=float)
+    parser.add_argument("--optimizer", type=str, default="adamw", choices=["adamw", "sgd"])
+    parser.add_argument("--momentum", type=float, default=0.9)
+    parser.add_argument("--warmup_epochs", type=int, default=5)
+    parser.add_argument("--start_epoch", type=int, default=0)
+    parser.add_argument("--save_freq", type=int, default=10)
+    parser.add_argument("--val_freq", type=int, default=5)
+    parser.add_argument("--resume", default="", type=str)
+    parser.add_argument("--gpu", default=0, type=int)
+
+    # Augmentation
+    parser.add_argument("--aug_scale_coords", action="store_true", default=False)
+    parser.add_argument("--aug_rotate", action="store_true", default=False)
+    parser.add_argument("--aug_jitter_coords", action="store_true", default=False)
+    parser.add_argument("--aug_shift_coords", action="store_true", default=False)
+    parser.add_argument("--aug_flip", action="store_true", default=False)
+    parser.add_argument("--aug_mask_feats", action="store_true", default=False)
+    parser.add_argument("--aug_jitter_length", action="store_true", default=False)
+    parser.add_argument("--aug_elasticate", action="store_true", default=False)
+    parser.add_argument("--aug_drop_tree", action="store_true", default=False)
+    parser.add_argument("--aug_skip_parent_node", action="store_true", default=False)
+    parser.add_argument("--aug_swap_sibling_subtrees", action="store_true", default=False)
+
+    # Persistence augmentation
+    parser.add_argument("--use_persistence_aug", action="store_true", default=False)
+    parser.add_argument("--pers_translation_scale", type=float, default=0.05)
+    parser.add_argument("--pers_noise_scale", type=float, default=0.02)
+    parser.add_argument("--pers_sigma_min", type=float, default=12.0)
+    parser.add_argument("--pers_sigma_max", type=float, default=20.0)
+    parser.add_argument("--sigma_px", type=float, default=16.0)
+
+    # Evaluation datasets
+    parser.add_argument("--eval_jm", action="store_true", default=False)
+    parser.add_argument("--eval_act", action="store_true", default=False)
+    parser.add_argument("--eval_neuron7", action="store_true", default=False)
+    parser.add_argument("--eval_m1_cell", action="store_true", default=False)
+    parser.add_argument("--eval_m1_region", action="store_true", default=False)
+    parser.add_argument("--eval_swc_glia", action="store_true", default=False)
+
+    # KNN evaluation
+    parser.add_argument("--use_knn_eval", action="store_true", default=False)
+    parser.add_argument("--knn_k", type=int, default=20)
+    parser.add_argument("--knn_t", type=float, default=0.5)
+    parser.add_argument("--knn_fusion", type=str, default="concat",
+                        choices=["concat", "add", "tree_only", "image_only"])
+    parser.add_argument("--use_sklearn_knn", action="store_true", default=False)
+
+    parser.add_argument("--cache_images", action="store_true", default=True)
+    parser.add_argument("--debug", action="store_true", default=False)
+
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    # Setup work directory
+    args.work_dir = f"{args.work_dir}/{args.exp_name}"
+    if not os.path.exists(args.work_dir):
+        os.makedirs(args.work_dir)
+
+    # Logger
+    timestamp = time.strftime("%Y%m%d_%H%M%S", time.localtime())
+    if args.debug:
+        log_file = None
+        args.save_freq = 10000
+        args.val_freq = 1
+    else:
+        log_file = f"{args.work_dir}/train_{timestamp}.log"
+    logger = get_root_logger(log_file=log_file, log_level="INFO")
+
+    logger.info("=" * 60)
+    logger.info("GraPHFormer TRAINING")
+    logger.info(f"Tree Encoder: {args.tree_model}")
+    logger.info(f"Image Encoder: {args.image_encoder}")
+    logger.info(f"Embedding Dimension: {args.embed_dim}")
+    logger.info(f"Temperature: {args.temperature}")
+    logger.info("=" * 60)
+    logger.info(json.dumps(vars(args), indent=4, sort_keys=True))
+
+    device = torch.device("cuda")
+
+    # Create model
+    logger.info("=> Creating model...")
+    model = CLIPModel(args).to(device)
+    logger.info(model)
+
+    # Optimizer
+    if args.optimizer == "sgd":
+        optimizer = torch.optim.SGD(
+            model.parameters(),
+            lr=args.lr,
+            weight_decay=args.wd,
+            momentum=args.momentum
+        )
+    else:
+        optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=args.lr,
+            weight_decay=args.wd,
+            betas=(0.9, 0.98),
+            eps=1e-6,
+        )
+
+    # Resume from checkpoint
+    if args.resume:
+        if os.path.isfile(args.resume):
+            logger.info(f"=> Loading checkpoint '{args.resume}'")
+            checkpoint = torch.load(args.resume, map_location=f"cuda:{args.gpu}")
+            args.start_epoch = checkpoint["epoch"]
+            model.load_state_dict(checkpoint["state_dict"])
+            optimizer.load_state_dict(checkpoint["optimizer"])
+            logger.info(f"=> Loaded checkpoint (epoch {checkpoint['epoch']})")
+
+    cudnn.benchmark = True
+
+    # Build augmentations
+    aug_switchs = [
+        False,
+        args.aug_scale_coords,
+        args.aug_rotate,
+        args.aug_jitter_coords,
+        args.aug_shift_coords,
+        args.aug_flip,
+        args.aug_mask_feats,
+        args.aug_jitter_length,
+        args.aug_elasticate,
+    ]
+    aug_fns = [
+        None,
+        RandomScaleCoords(p=0.2),
+        RandomRotate(p=0.5),
+        RandomJitter(p=0.2),
+        RandomShift(p=0.2),
+        RandomFlip(p=1),
+        RandomMaskFeats(p=0.2),
+        RandomJitterLength(p=0.2),
+        RandomElasticate(p=0.2),
+    ]
+    feat_augs = [aug_fns[i] for i in range(len(aug_switchs)) if aug_switchs[i] and aug_fns[i] is not None]
+    feat_augs = Compose(feat_augs) if feat_augs else None
+
+    topo_aug_switchs = [
+        args.aug_drop_tree,
+        args.aug_skip_parent_node,
+        args.aug_swap_sibling_subtrees,
+    ]
+    topo_aug_fns = [
+        RandomDropSubTrees(probs=[0.05], max_cnt=5),
+        RandomSkipParentNode(probs=[0.05], max_cnt=10),
+        RandomSwapSiblingSubTrees(probs=[0.05], max_cnt=10),
+    ]
+    topo_augs = [topo_aug_fns[i] for i in range(len(topo_aug_switchs)) if topo_aug_switchs[i]]
+    topo_augs = Compose(topo_augs) if topo_augs else None
+
+    # Persistence augmentation
+    persistence_aug = None
+    if args.use_persistence_aug:
+        persistence_aug = CombinedPersistenceAugmentation(
+            translation_scale=args.pers_translation_scale,
+            noise_scale=args.pers_noise_scale,
+            sigma_min=args.pers_sigma_min,
+            sigma_max=args.pers_sigma_max,
+        )
+
+    # Create training dataset
+    use_full_phase = args.dataset in ["all_wo_others", "all_with_neuron7", "neuron7", "ACT"]
+
+    trainset = NeuronTreeDataset(
+        phase="full" if use_full_phase else "train",
+        dataset=args.dataset,
+        label_dict=LABEL_DICT[args.dataset],
+        data_dir=args.data_dir,
+        topology_transformations=topo_augs,
+        attribute_transformations=feat_augs,
+        input_features=args.input_features,
+        use_images=True,
+        image_size=args.image_size,
+        cache_images=args.cache_images,
+        persistence_augmentation=persistence_aug,
+        sigma_px=args.sigma_px,
+    )
+
+    collate_fn = get_collate_fn(device, use_images=True)
+
+    train_loader = DataLoader(
+        dataset=trainset,
+        batch_size=args.batch_size,
+        collate_fn=collate_fn,
+        shuffle=True,
+        drop_last=True,
+        num_workers=6,
+        pin_memory=True,
+        persistent_workers=True,
+    )
+
+    # Evaluation datasets
+    eval_datasets = []
+    eval_loaders = []
+    eval_memory_loaders = []
+
+    # BIL (always evaluated)
+    bil_testset = create_eval_dataset("test", "bil_6_classes", args)
+    bil_test_loader = DataLoader(
+        dataset=bil_testset,
+        batch_size=args.batch_size,
+        collate_fn=collate_fn,
+        shuffle=False,
+        num_workers=4,
+        pin_memory=True,
+    )
+    eval_datasets.append("BIL")
+    eval_loaders.append(bil_test_loader)
+
+    if args.use_knn_eval:
+        bil_memory = create_eval_dataset("train", "bil_6_classes", args)
+        bil_memory_loader = DataLoader(
+            dataset=bil_memory,
+            batch_size=args.batch_size,
+            collate_fn=collate_fn,
+            shuffle=False,
+            num_workers=4,
+            pin_memory=True,
+        )
+        eval_memory_loaders.append(bil_memory_loader)
+
+    # JM
+    if args.eval_jm:
+        jm_testset = create_eval_dataset("test", "JM", args)
+        jm_test_loader = DataLoader(
+            dataset=jm_testset,
+            batch_size=args.batch_size,
+            collate_fn=collate_fn,
+            shuffle=False,
+            num_workers=4,
+            pin_memory=True,
+        )
+        eval_datasets.append("JM")
+        eval_loaders.append(jm_test_loader)
+
+        if args.use_knn_eval:
+            jm_memory = create_eval_dataset("train", "JM", args)
+            jm_memory_loader = DataLoader(
+                dataset=jm_memory,
+                batch_size=args.batch_size,
+                collate_fn=collate_fn,
+                shuffle=False,
+                num_workers=4,
+                pin_memory=True,
+            )
+            eval_memory_loaders.append(jm_memory_loader)
+
+    # ACT
+    if args.eval_act:
+        act_testset = create_eval_dataset("test", "ACT", args)
+        act_test_loader = DataLoader(
+            dataset=act_testset,
+            batch_size=args.batch_size,
+            collate_fn=collate_fn,
+            shuffle=False,
+            num_workers=4,
+            pin_memory=True,
+        )
+        eval_datasets.append("ACT")
+        eval_loaders.append(act_test_loader)
+
+        if args.use_knn_eval:
+            act_memory = create_eval_dataset("train", "ACT", args)
+            act_memory_loader = DataLoader(
+                dataset=act_memory,
+                batch_size=args.batch_size,
+                collate_fn=collate_fn,
+                shuffle=False,
+                num_workers=4,
+                pin_memory=True,
+            )
+            eval_memory_loaders.append(act_memory_loader)
+
+    # Neuron7
+    if args.eval_neuron7:
+        neuron7_testset = create_eval_dataset("test", "neuron7", args)
+        neuron7_test_loader = DataLoader(
+            dataset=neuron7_testset,
+            batch_size=args.batch_size,
+            collate_fn=collate_fn,
+            shuffle=False,
+            num_workers=4,
+            pin_memory=True,
+        )
+        eval_datasets.append("neuron7")
+        eval_loaders.append(neuron7_test_loader)
+
+        if args.use_knn_eval:
+            neuron7_memory = create_eval_dataset("train", "neuron7", args)
+            neuron7_memory_loader = DataLoader(
+                dataset=neuron7_memory,
+                batch_size=args.batch_size,
+                collate_fn=collate_fn,
+                shuffle=False,
+                num_workers=4,
+                pin_memory=True,
+            )
+            eval_memory_loaders.append(neuron7_memory_loader)
+
+    # M1_EXC_cell
+    if args.eval_m1_cell:
+        m1_cell_testset = create_eval_dataset("test", "m1_exc_cell", args)
+        m1_cell_test_loader = DataLoader(
+            dataset=m1_cell_testset,
+            batch_size=args.batch_size,
+            collate_fn=collate_fn,
+            shuffle=False,
+            num_workers=4,
+            pin_memory=True,
+        )
+        eval_datasets.append("m1_exc_cell")
+        eval_loaders.append(m1_cell_test_loader)
+
+        if args.use_knn_eval:
+            m1_cell_memory = create_eval_dataset("train", "m1_exc_cell", args)
+            m1_cell_memory_loader = DataLoader(
+                dataset=m1_cell_memory,
+                batch_size=args.batch_size,
+                collate_fn=collate_fn,
+                shuffle=False,
+                num_workers=4,
+                pin_memory=True,
+            )
+            eval_memory_loaders.append(m1_cell_memory_loader)
+
+    # M1_EXC_region
+    if args.eval_m1_region:
+        m1_region_testset = create_eval_dataset("test", "m1_exc_region", args)
+        m1_region_test_loader = DataLoader(
+            dataset=m1_region_testset,
+            batch_size=args.batch_size,
+            collate_fn=collate_fn,
+            shuffle=False,
+            num_workers=4,
+            pin_memory=True,
+        )
+        eval_datasets.append("m1_exc_region")
+        eval_loaders.append(m1_region_test_loader)
+
+        if args.use_knn_eval:
+            m1_region_memory = create_eval_dataset("train", "m1_exc_region", args)
+            m1_region_memory_loader = DataLoader(
+                dataset=m1_region_memory,
+                batch_size=args.batch_size,
+                collate_fn=collate_fn,
+                shuffle=False,
+                num_workers=4,
+                pin_memory=True,
+            )
+            eval_memory_loaders.append(m1_region_memory_loader)
+
+    # swc_glia
+    if args.eval_swc_glia:
+        swc_glia_testset = create_eval_dataset("test", "swc_glia_filtered_1000", args)
+        swc_glia_test_loader = DataLoader(
+            dataset=swc_glia_testset,
+            batch_size=args.batch_size,
+            collate_fn=collate_fn,
+            shuffle=False,
+            num_workers=4,
+            pin_memory=True,
+        )
+        eval_datasets.append("swc_glia")
+        eval_loaders.append(swc_glia_test_loader)
+
+        if args.use_knn_eval:
+            swc_glia_memory = create_eval_dataset("train", "swc_glia_filtered_1000", args)
+            swc_glia_memory_loader = DataLoader(
+                dataset=swc_glia_memory,
+                batch_size=args.batch_size,
+                collate_fn=collate_fn,
+                shuffle=False,
+                num_workers=4,
+                pin_memory=True,
+            )
+            eval_memory_loaders.append(swc_glia_memory_loader)
+
+    # Learning rate scheduler
+    def lr_schedule(epoch):
+        if epoch < args.warmup_epochs:
+            return (epoch + 1) / args.warmup_epochs
+        else:
+            progress = (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)
+            cosine_decay = 0.5 * (1 + torch.cos(torch.tensor(progress * 3.14159)))
+            min_lr_factor = 1e-6 / args.lr
+            return min_lr_factor + (1 - min_lr_factor) * cosine_decay
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_schedule)
+
+    # Training loop
+    best_metrics = {dataset: {"recall@5": 0.0, "epoch": 0} for dataset in eval_datasets}
+    total_iters = len(train_loader) * args.epochs
+    current_iter = 0
+    start_time = time.time()
+
+    logger.info("=> Starting training...")
+    for epoch in range(args.start_epoch + 1, args.epochs + 1):
+        model.train()
+        epoch_loss = 0.0
+
+        for step, batch in enumerate(train_loader):
+            try:
+                batch = batch.to(device)
+                loss = model(batch)
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+                epoch_loss += loss.item()
+                current_iter += 1
+
+                if step % 10 == 0:
+                    current_time = time.time()
+                    elapsed = current_time - start_time
+
+                    log_str = (
+                        f"Epoch {epoch:03d} | Step {step:03d}/{len(train_loader)} | "
+                        f"Loss {loss.item():.4f} | "
+                        f"LR {optimizer.param_groups[0]['lr']:.6f} | "
+                        f"Elapsed {str(datetime.timedelta(seconds=int(elapsed)))}"
+                    )
+                    logger.info(log_str)
+            except Exception as e:
+                logger.info(f"Error in step {step}: {e}")
+                continue
+
+        scheduler.step()
+        avg_loss = epoch_loss / len(train_loader)
+        logger.info(f"Epoch {epoch:03d} | Avg Loss: {avg_loss:.4f}")
+
+        # Evaluation
+        if epoch % args.val_freq == 0 and args.use_knn_eval:
+            logger.info("=> Evaluating...")
+
+            fusion_modes = ['concat', 'add', 'tree_only', 'image_only']
+
+            for idx, (dataset_name, test_loader) in enumerate(zip(eval_datasets, eval_loaders)):
+                memory_loader = eval_memory_loaders[idx]
+                num_classes = len(test_loader.dataset.classes)
+
+                # Use k=5 for JM, otherwise use args.knn_k
+                k = 5 if dataset_name == "JM" else args.knn_k
+
+                logger.info(f"\n  === {dataset_name} Dataset ===")
+
+                fusion_results = {}
+                best_fusion_acc = 0
+                best_fusion_mode = None
+
+                for fusion_mode in fusion_modes:
+                    if args.use_sklearn_knn:
+                        knn_acc = evaluate_sklearn_knn(
+                            model, memory_loader, test_loader, device,
+                            knn_k=k, fusion=fusion_mode
+                        )
+                        knn_acc = knn_acc * 100
+                    else:
+                        knn_acc = evaluate_knn(
+                            model, memory_loader, test_loader, device,
+                            num_classes, k, args.knn_t, fusion=fusion_mode
+                        )
+
+                    fusion_results[fusion_mode] = knn_acc
+
+                    if knn_acc > best_fusion_acc:
+                        best_fusion_acc = knn_acc
+                        best_fusion_mode = fusion_mode
+
+                    logger.info(f"    {fusion_mode:12s}: {knn_acc:.2f}%")
+
+                logger.info(f"    {'BEST':12s}: {best_fusion_mode} ({best_fusion_acc:.2f}%)")
+
+                # Save best checkpoint based on higher of concat or add
+                concat_acc = fusion_results['concat']
+                add_acc = fusion_results['add']
+                primary_acc = max(concat_acc, add_acc)
+                primary_fusion = 'add' if add_acc > concat_acc else 'concat'
+
+                logger.info(f"    Best concat : {concat_acc:.2f}%")
+                logger.info(f"    Best add    : {add_acc:.2f}%")
+                logger.info(f"    Selected    : {primary_fusion} ({primary_acc:.2f}%)")
+
+                if primary_acc > best_metrics[dataset_name]["recall@5"]:
+                    best_metrics[dataset_name]["recall@5"] = primary_acc
+                    best_metrics[dataset_name]["epoch"] = epoch
+                    best_metrics[dataset_name]["fusion_mode"] = primary_fusion
+                    best_metrics[dataset_name]["best_concat"] = concat_acc
+                    best_metrics[dataset_name]["best_add"] = add_acc
+
+                    checkpoint_path = f"{args.work_dir}/best_{dataset_name}_epoch_{epoch}.pth"
+                    save_checkpoint(
+                        {
+                            "epoch": epoch,
+                            "state_dict": model.state_dict(),
+                            "optimizer": optimizer.state_dict(),
+                            "knn_accuracy": primary_acc,
+                            "primary_fusion_mode": primary_fusion,
+                            "concat_accuracy": concat_acc,
+                            "add_accuracy": add_acc,
+                            "fusion_results": fusion_results,
+                            "dataset": dataset_name,
+                        },
+                        is_best=True,
+                        filename=checkpoint_path,
+                    )
+                    logger.info(f"  Saved new best for {dataset_name}: {checkpoint_path}")
+
+                logger.info(
+                    f"  Best {dataset_name} ({best_metrics[dataset_name].get('fusion_mode', primary_fusion)}): "
+                    f"{best_metrics[dataset_name]['recall@5']:.2f}% at epoch {best_metrics[dataset_name]['epoch']} "
+                    f"[concat: {best_metrics[dataset_name].get('best_concat', concat_acc):.2f}%, "
+                    f"add: {best_metrics[dataset_name].get('best_add', add_acc):.2f}%]"
+                )
+
+        # Save periodic checkpoint
+        if epoch % args.save_freq == 0:
+            checkpoint_path = f"{args.work_dir}/epoch_{epoch}.pth"
+            save_checkpoint(
+                {
+                    "epoch": epoch,
+                    "state_dict": model.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                },
+                is_best=False,
+                filename=checkpoint_path,
+            )
+            logger.info(f"Saved checkpoint: {checkpoint_path}")
+
+    logger.info("Training complete!")
+    if args.use_knn_eval:
+        logger.info("Best results:")
+        for dataset_name in eval_datasets:
+            logger.info(
+                f"  {dataset_name}: {best_metrics[dataset_name]['recall@5']:.2f}% "
+                f"at epoch {best_metrics[dataset_name]['epoch']}"
+            )