Baseline/architecture_ablation.py

# EMMA Architecture Ablation Study: LTC vs LSTM vs GRU vs Transformer
# Critical Ablation A.1 for Paper

import os
import csv
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from ncps.torch import LTC

# Set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

Nloop = 0


class Custom_Pendulum_Loss(nn.Module):
    """Physics-informed loss identical to run-45 pipeline."""
    
    def __init__(self, labels, logits, omega):
        super().__init__()
        self.y_true = labels
        self.y_pred = logits
        self.y_omega = omega
    
    def forward(self):
        dev = self.y_pred.device
        T, B, _ = self.y_pred.shape
        
        maxChange = 95.0
        getp = lambda k: self.y_pred[:,:,k]
        
        alpha_nominal = 0.45
        beta_nominal = 0.05
        gamma_nominal = 100.0
        
        alpha = (1 + (0.5 - getp(0)) * maxChange / 100.0) * alpha_nominal
        beta = (1 + (0.5 - getp(1)) * maxChange / 100.0) * beta_nominal
        gamma = (1 + (0.5 - getp(2)) * maxChange / 100.0) * gamma_nominal
        
        L = alpha
        tau = beta
        g = torch.tensor(9.81, device=dev)
        
        thetaVal = self.y_true[:,:,0]
        omegaVal = self.y_omega[:,:,0]
        theta = thetaVal.clone().unsqueeze(2)
        omega = omegaVal.clone().unsqueeze(2)
        
        limitLoop = min(500, T)
        tau_dt = 0.03
        
        for i in range(1, limitLoop):
            y1 = theta[:,:,i-1] + omega[:,:,i-1]*tau_dt
            y0 = omega[:,:,i-1] + (-torch.mul(tau,omega[:,:,i-1]) - torch.mul(torch.div(g,L.clamp(min=0.0001)),torch.sin(theta[:,:,i-1])))*tau_dt
            theta = torch.cat([theta, y1.unsqueeze(2)],dim=2)
            omega = torch.cat([omega, y0.unsqueeze(2)],dim=2)
        
        loss_Cal_theta = gamma * 0.01
        loss_Cal_omega = gamma * 0.005
        
        mse_loss = torch.abs(torch.sum(torch.square(self.y_true[:,:,0:limitLoop]-theta)/limitLoop, dim=2)-loss_Cal_theta) + \
                   torch.abs(torch.sum(torch.square(self.y_omega[:,:,0:limitLoop]-omega)/limitLoop, dim=2)-loss_Cal_omega)
        
        param_penalty = 0.0
        param_penalty += 10.0 * torch.mean(torch.relu(-alpha))
        param_penalty += 10.0 * torch.mean(torch.relu(-beta))
        param_penalty += 10.0 * torch.mean(torch.relu(-gamma))
        param_penalty += 2.0 * torch.mean(torch.relu(alpha - 2.0))
        param_penalty += 2.0 * torch.mean(torch.relu(beta - 1.0))
        param_penalty += 1.0 * torch.mean(torch.relu(gamma - 500.0))
        
        total_loss = mse_loss + 0.001 * param_penalty
        
        self.L = L
        self.tau = tau
        return total_loss


def cut_in_sequences(x, y, seq_len, inc=1):
    sequences_x, sequences_y = [], []
    for s in range(0, x.shape[0] - seq_len, inc):
        start, end = s, s + seq_len
        sequences_x.append(x[start:end])
        sequences_y.append(y[start:end])
    return np.stack(sequences_x, axis=1), np.stack(sequences_y, axis=1)


class PendulumData:
    def __init__(self, seq_len=16, data_dir="data"):
        print(f"Loading data from {data_dir}...")
        theta_data = np.loadtxt(os.path.join(data_dir, "thetaData.txt"))
        omega_data = np.loadtxt(os.path.join(data_dir, "omegaData.txt"))
        theta_traj = theta_data.T
        omega_traj = omega_data.T
        global Nloop
        Nloop = theta_traj.shape[1]
        
        train_x, train_y = cut_in_sequences(theta_traj, theta_traj, seq_len)
        train_omega, train_omega_y = cut_in_sequences(omega_traj, omega_traj, seq_len)
        
        self.train_x = torch.tensor(train_x, dtype=torch.float32)
        self.train_y = torch.tensor(train_y, dtype=torch.float32)
        self.train_omega = torch.tensor(train_omega, dtype=torch.float32)
        self.train_omega_y = torch.tensor(train_omega_y, dtype=torch.float32)
        
        print(f"Training sequences: {self.train_x.shape[1]}")
    
    def iterate_train(self, batch_size=32):
        total_seqs = self.train_x.shape[1]
        total_batches = max(1, total_seqs // batch_size)
        for i in range(total_batches):
            start = i * batch_size
            end = start + batch_size
            yield (self.train_x[:, start:end], self.train_y[:, start:end], 
                   self.train_omega[:, start:end], self.train_omega_y[:, start:end])


class ArchitectureModel(nn.Module):
    """Unified model supporting LTC, LSTM, GRU, and Transformer."""
    
    def __init__(self, model_type="ltc", model_size=64, learning_rate=0.0003):
        super().__init__()
        self.model_type = model_type.lower()
        self.model_size = model_size
        
        input_size = Nloop if Nloop > 0 else 100
        print(f"Building {model_type.upper()} model with {model_size} units...")
        
        if self.model_type == "ltc":
            self.wm = LTC(
                input_size=input_size,
                units=model_size,
                return_sequences=True,
                batch_first=False,
                mixed_memory=False,
                ode_unfolds=6,
                epsilon=1e-8
            )
            self.rnn = self.wm
            learning_rate = 0.005
        elif self.model_type == "lstm":
            self.rnn = nn.LSTM(input_size, model_size, batch_first=False)
        elif self.model_type == "gru":
            self.rnn = nn.GRU(input_size, model_size, batch_first=False)
        elif self.model_type == "transformer":
            nhead = 4
            padded_input_size = ((input_size + nhead - 1) // nhead) * nhead
            if input_size != padded_input_size:
                self.input_projection = nn.Linear(input_size, padded_input_size)
            else:
                self.input_projection = None
            
            encoder_layer = nn.TransformerEncoderLayer(
                d_model=padded_input_size, 
                nhead=nhead, 
                dim_feedforward=model_size,
                batch_first=False
            )
            self.rnn = nn.TransformerEncoder(encoder_layer, num_layers=2)
            self.transformer_out = nn.Linear(padded_input_size, model_size)
        else:
            raise ValueError(f"Unknown model type: {model_type}")
        
        self.dense = nn.Linear(model_size, 3)
        self.sigmoid = nn.Sigmoid()
        
        self.optimizer = optim.AdamW(self.parameters(), lr=learning_rate, 
                                    weight_decay=1e-4, betas=(0.9, 0.999), eps=1e-8)
        self.to(device)
        
        self.scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
            self.optimizer, T_0=10, T_mult=2, eta_min=1e-6
        )
    
    def forward(self, x):
        if self.model_type == "ltc":
            out, _ = self.rnn(x)
        elif self.model_type == "transformer":
            if hasattr(self, 'input_projection') and self.input_projection is not None:
                T, B, _ = x.shape
                x = self.input_projection(x.reshape(T*B, -1)).reshape(T, B, -1)
            out = self.rnn(x)
            out = self.transformer_out(out)
        else:  # LSTM, GRU
            out, _ = self.rnn(x)
        
        T, B, H = out.shape
        y = self.sigmoid(self.dense(out.reshape(T*B, H))).reshape(T, B, 3)
        return y
    
    def compute_loss(self, y_pred, target_y, omega):
        loss_fn = Custom_Pendulum_Loss(target_y, y_pred, omega)
        return loss_fn.forward()


def run_architecture_ablation():
    """Run architecture comparison: LTC vs LSTM vs GRU vs Transformer."""
    
    # Get absolute path
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # Configuration
    pendulum_45_datasets = [
        "Pendulum-EMMA/45_v1",
        "Pendulum-EMMA/45_v2",
        "Pendulum-EMMA/45_v3",
        "Pendulum-EMMA/45_v4",
        "Pendulum-EMMA/45_v5"
    ]
    
    # Architecture comparison (all with 64 hidden units for fair comparison)
    architectures = ["ltc", "lstm", "gru", "transformer"]
    
    # Training parameters
    seq_len = 16
    batch_size = 2
    num_epochs = 40
    learning_rate = 0.0003
    
    results = []
    
    print("=" * 70)
    print("EMMA ARCHITECTURE ABLATION STUDY")
    print("Comparison: LTC vs LSTM vs GRU vs Transformer (NODE)")
    print("=" * 70)
    
    total_experiments = len(architectures) * len(pendulum_45_datasets)
    experiment_num = 0
    
    for dataset_path in pendulum_45_datasets:
        dataset_name = os.path.basename(dataset_path)
        print(f"\n{'='*70}")
        print(f"DATASET: {dataset_name}")
        print(f"{'='*70}")
        
        # Load data
        data_dir = os.path.join(script_dir, dataset_path, "data")
        if not os.path.exists(data_dir):
            print(f"[WARNING] Data directory not found: {data_dir}")
            continue
            
        dataset = PendulumData(seq_len=seq_len, data_dir=data_dir)
        
        for arch_name in architectures:
            experiment_num += 1
            print(f"\n[{experiment_num}/{total_experiments}] Architecture: {arch_name.upper()}, Dataset: {dataset_name}")
            
            # Set seed
            seed = int(dataset_name.split('_v')[-1])
            np.random.seed(seed)
            torch.manual_seed(seed)
            
            # Create model
            model = ArchitectureModel(model_type=arch_name, model_size=64, learning_rate=learning_rate).to(device)
            optimizer = model.optimizer
            scheduler = model.scheduler
            
            # Training loop with convergence tracking
            best_loss = float('inf')
            convergence_epoch = num_epochs
            loss_history = []
            import time
            start_time = time.time()
            patience = 5
            patience_counter = 0
            convergence_threshold = 1e-4  # Loss improvement threshold
            
            for epoch in range(num_epochs):
                model.train()
                epoch_loss = 0.0
                batch_count = 0
                
                for batch_x, batch_y, batch_omega, batch_omega_y in dataset.iterate_train(batch_size=batch_size):
                    batch_x = batch_x.to(device)
                    batch_y = batch_y.to(device)
                    batch_omega = batch_omega.to(device)
                    
                    optimizer.zero_grad()
                    predicted_params = model(batch_x)
                    loss_mat = model.compute_loss(predicted_params, batch_y, batch_omega)
                    loss = loss_mat.mean()
                    
                    if torch.isnan(loss):
                        continue
                    
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                    optimizer.step()
                    
                    epoch_loss += loss.item()
                    batch_count += 1
                
                if batch_count > 0:
                    avg_loss = epoch_loss / batch_count
                    loss_history.append(avg_loss)
                    scheduler.step()
                    
                    # Track convergence (when loss improvement is minimal)
                    if avg_loss < best_loss - convergence_threshold:
                        best_loss = avg_loss
                        patience_counter = 0
                        convergence_epoch = epoch + 1
                    else:
                        patience_counter += 1
                    
                    if (epoch + 1) % 10 == 0:
                        print(f'  Epoch {epoch+1}/{num_epochs}, Loss: {avg_loss:.6f}')
                    
                    # Early stopping if converged
                    if patience_counter >= patience and epoch >= 10:
                        print(f'  Converged at epoch {convergence_epoch}')
                        break
            
            training_time = time.time() - start_time
            final_loss = loss_history[-1] if loss_history else best_loss
            
            # Evaluate
            model.eval()
            with torch.no_grad():
                sample_batch = next(iter(dataset.iterate_train(batch_size=1)))
                sample_x, sample_y, sample_omega, sample_omega_y = sample_batch
                
                sample_x = sample_x.to(device)
                predicted_params = model(sample_x)
                
                maxChange = 95.0
                getp = lambda k: predicted_params[:,:,k].mean()
                
                L = ((1 + (0.5 - getp(0)) * maxChange / 100.0) * 0.45).item()
                tau = ((1 + (0.5 - getp(1)) * maxChange / 100.0) * 0.05).item()
                
                # Calculate parameter recovery accuracy (how close to ground truth)
                L_gt = 0.45  # Ground truth for 45° pendulum
                tau_gt = 0.05  # Ground truth damping
                L_error = abs(L - L_gt) / L_gt * 100  # Percentage error
                tau_error = abs(tau - tau_gt) / tau_gt * 100
                
                print(f"  L: {L:.6f} m (error: {L_error:.2f}%), tau: {tau:.6f} 1/s (error: {tau_error:.2f}%)")
                print(f"  Time: {training_time:.2f}s, Converged: epoch {convergence_epoch}, Final loss: {final_loss:.6f}")
                
                results.append({
                    'dataset': dataset_name,
                    'architecture': arch_name.upper(),
                    'hidden_units': 64,
                    'video_number': seed,
                    'best_loss': best_loss,
                    'final_loss': final_loss,
                    'training_time_s': training_time,
                    'convergence_epoch': convergence_epoch,
                    'L_estimated': L,
                    'tau_estimated': tau,
                    'L_error_percent': L_error,
                    'tau_error_percent': tau_error
                })
    
    # Save results
    results_file = os.path.join(script_dir, "architecture_ablation_results.csv")
    print(f"\n{'='*70}")
    print("SAVING RESULTS")
    print(f"{'='*70}")
    
    with open(results_file, 'w', newline='') as f:
        if results:
            writer = csv.DictWriter(f, fieldnames=results[0].keys())
            writer.writeheader()
            writer.writerows(results)
    
    print(f"Results saved to: {results_file}")
    
    # Print summary with key metrics
    print(f"\n{'='*70}")
    print("RESULTS SUMMARY - KEY METRICS FOR PAPER")
    print(f"{'='*70}")
    
    for arch_name in architectures:
        arch_results = [r for r in results if r['architecture'] == arch_name.upper()]
        if arch_results:
            L_values = [r['L_estimated'] for r in arch_results]
            tau_values = [r['tau_estimated'] for r in arch_results]
            L_errors = [r['L_error_percent'] for r in arch_results]
            tau_errors = [r['tau_error_percent'] for r in arch_results]
            time_values = [r['training_time_s'] for r in arch_results]
            conv_epochs = [r['convergence_epoch'] for r in arch_results]
            final_losses = [r['final_loss'] for r in arch_results]
            
            print(f"\n{arch_name.upper()}:")
            print(f"  Parameter Recovery Accuracy:")
            print(f"    L error:   {np.mean(L_errors):.2f}% ± {np.std(L_errors):.2f}%")
            print(f"    tau error: {np.mean(tau_errors):.2f}% ± {np.std(tau_errors):.2f}%")
            print(f"  Convergence Speed:")
            print(f"    Epochs to converge: {np.mean(conv_epochs):.1f} ± {np.std(conv_epochs):.1f}")
            print(f"  Stability (lower is better):")
            print(f"    L std:     {np.std(L_values):.6f} m")
            print(f"    tau std:   {np.std(tau_values):.6f} 1/s")
            print(f"    Loss std:  {np.std(final_losses):.6f}")
            print(f"  Training Time:")
            print(f"    {np.mean(time_values):.2f} ± {np.std(time_values):.2f} s")


if __name__ == "__main__":
    run_architecture_ablation()