plot_tsp_comparison.py

#!/usr/bin/env python3
"""
Plot TSP signals for all Alanine reference files
Compares raw (magnitude) vs phase-corrected (absorption) spectra
"""

import nmrglue as ng
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import os


def process_absorption(dic, data):
    """Process spectrum to absorption mode with phase correction"""
    real = data[0]
    imag = data[1] if len(data) > 1 else np.zeros_like(real)
    
    # Get phase correction parameters
    phc0 = float(dic.get('$PHC0', [0])[0])
    phc1 = float(dic.get('$PHC1', [0])[0])
    
    # Apply phase correction
    n = len(real)
    complex_data = real + 1j * imag
    phase = np.radians(phc0) + np.radians(phc1) * (np.arange(n) - n/2) / n
    phased = complex_data * np.exp(-1j * phase)
    
    return np.real(phased), np.imag(phased), np.sqrt(real**2 + imag**2)


def get_ppm_scale(dic, n_points):
    """Calculate ppm scale from dictionary"""
    sfo1 = float(dic['$SFO1'][0])
    o1_hz = float(dic['$O1'][0])
    sw_hz = float(dic['$SWH'][0])
    o1_ppm = o1_hz / sfo1
    sw_ppm = sw_hz / sfo1
    return np.linspace(o1_ppm + sw_ppm/2, o1_ppm - sw_ppm/2, n_points)


def plot_tsp_comparison():
    """Plot TSP signals for all Alanine files"""
    
    base_dir = "raw_data/Reference_Raw_Date_JCAMP-DX/Alanine-Reference"
    files = [10, 20, 30, 40, 50, 60, 70]
    
    # Prepare figure with subplots
    fig, axes = plt.subplots(2, 1, figsize=(14, 10))
    
    # Colors for different files
    colors = plt.cm.tab10(np.linspace(0, 1, len(files)))
    
    # Store data for summary
    tsp_data = []
    
    for i, fileno in enumerate(files):
        filepath = os.path.join(base_dir, f"{fileno}.dx")
        if not os.path.exists(filepath):
            print(f"File {fileno} not found, skipping")
            continue
        
        # Load data
        dic, data = ng.jcampdx.read(filepath)
        absorption, dispersion, magnitude = process_absorption(dic, data)
        ppm = get_ppm_scale(dic, len(absorption))
        
        # Find TSP peak position
        mask_tsp = (ppm >= -0.5) & (ppm <= 0.5)
        tsp_idx = np.argmax(magnitude[mask_tsp])
        tsp_pos = ppm[mask_tsp][tsp_idx]
        
        # Shift to center TSP at 0
        ppm_shifted = ppm - tsp_pos
        
        # Extract TSP region (-0.5 to 0.5 ppm around TSP)
        mask_region = (ppm_shifted >= -0.5) & (ppm_shifted <= 0.5)
        ppm_region = ppm_shifted[mask_region]
        mag_region = magnitude[mask_region]
        abs_region = absorption[mask_region]
        
        # Store data
        tsp_data.append({
            'fileno': fileno,
            'tsp_pos': tsp_pos,
            'mag_max': np.max(mag_region),
            'abs_max': np.max(abs_region),
            'mag_integral': np.trapz(mag_region, ppm_region),
            'abs_integral': np.trapz(abs_region, ppm_region)
        })
        
        # Plot raw magnitude
        axes[0].plot(ppm_region, mag_region, color=colors[i], 
                    label=f'File {fileno}', linewidth=1.5, alpha=0.8)
        
        # Plot phase-corrected absorption
        axes[1].plot(ppm_region, abs_region, color=colors[i],
                    label=f'File {fileno}', linewidth=1.5, alpha=0.8)
    
    # Configure plots
    axes[0].set_xlabel('Chemical Shift (ppm)', fontsize=12)
    axes[0].set_ylabel('Intensity (Magnitude)', fontsize=12)
    axes[0].set_title('TSP Region - Raw Magnitude Spectra', fontsize=14, fontweight='bold')
    axes[0].legend(loc='upper right', fontsize=9)
    axes[0].grid(True, alpha=0.3)
    axes[0].axvline(x=0, color='k', linestyle='--', alpha=0.5, label='TSP center')
    
    axes[1].set_xlabel('Chemical Shift (ppm)', fontsize=12)
    axes[1].set_ylabel('Intensity (Absorption)', fontsize=12)
    axes[1].set_title('TSP Region - Phase-Corrected Absorption Spectra', fontsize=14, fontweight='bold')
    axes[1].legend(loc='upper right', fontsize=9)
    axes[1].grid(True, alpha=0.3)
    axes[1].axvline(x=0, color='k', linestyle='--', alpha=0.5)
    
    plt.tight_layout()
    plt.savefig('tsp_comparison_alanine.png', dpi=150, bbox_inches='tight')
    print("Saved: tsp_comparison_alanine.png")
    
    # Create summary table
    print("\n" + "="*100)
    print("TSP SIGNAL SUMMARY - Alanine Reference Library")
    print("="*100)
    print(f"{'File':<8} {'TSP Pos':<10} {'Mag Max':<15} {'Abs Max':<15} {'Mag Integral':<18} {'Abs Integral':<18}")
    print("-"*100)
    
    for d in tsp_data:
        print(f"{d['fileno']:<8} {d['tsp_pos']:<10.4f} {d['mag_max']:<15.2e} {d['abs_max']:<15.2e} "
              f"{d['mag_integral']:<18.2e} {d['abs_integral']:<18.2e}")
    
    # Calculate Scale(TSP) relative to File 10
    print("\n" + "="*100)
    print("Scale(TSP) Calculations (relative to File 10)")
    print("="*100)
    
    ref_data = next(d for d in tsp_data if d['fileno'] == 10)
    
    print(f"{'File':<8} {'Scale(Mag)':<15} {'Scale(Abs)':<15} {'Scale(Mag Int)':<18} {'Scale(Abs Int)':<18}")
    print("-"*100)
    
    for d in tsp_data:
        if d['fileno'] != 10:
            scale_mag = d['mag_max'] / ref_data['mag_max']
            scale_abs = d['abs_max'] / ref_data['abs_max']
            scale_mag_int = d['mag_integral'] / ref_data['mag_integral']
            scale_abs_int = d['abs_integral'] / ref_data['abs_integral']
            
            print(f"{d['fileno']:<8} {scale_mag:<15.4f} {scale_abs:<15.4f} "
                  f"{scale_mag_int:<18.4f} {scale_abs_int:<18.4f}")
    
    print("\n" + "="*100)
    print("Comparison with Table S5 (for reference):")
    print("-"*100)
    table_s5 = {20: 0.9932, 30: 0.6607, 40: 0.8964, 50: 0.8992, 60: 0.9546, 70: 0.8852}
    print(f"{'File':<8} {'Our Scale':<15} {'Table S5':<15} {'Ratio':<10}")
    print("-"*100)
    for d in tsp_data:
        if d['fileno'] in table_s5:
            our_scale = d['mag_integral'] / ref_data['mag_integral']
            table_val = table_s5[d['fileno']]
            print(f"{d['fileno']:<8} {our_scale:<15.4f} {table_val:<15.4f} {our_scale/table_val:<10.4f}")
    
    print("="*100)
    
    plt.close()


if __name__ == "__main__":
    plot_tsp_comparison()