plot_tsp_comparison_full_range.py

#!/usr/bin/env python3
"""
Plot full-range spectra for all Alanine reference files
Shows magnitude spectra across entire ppm range
"""

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_full_range_spectra():
    """Plot full-range spectra 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 - full range and zoomed regions
    fig = plt.figure(figsize=(16, 12))
    gs = fig.add_gridspec(3, 2, height_ratios=[2, 1, 1])
    
    ax_full = fig.add_subplot(gs[0, :])  # Full spectrum spans both columns
    ax_tsp = fig.add_subplot(gs[1, 0])   # TSP region
    ax_methyl = fig.add_subplot(gs[1, 1])  # Methyl region
    ax_ch = fig.add_subplot(gs[2, 0])    # CH region
    ax_summary = fig.add_subplot(gs[2, 1])  # Summary table
    
    # 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
        
        # Store data
        tsp_region = (ppm_shifted >= -0.5) & (ppm_shifted <= 0.5)
        tsp_max = np.max(magnitude[tsp_region])
        
        tsp_data.append({
            'fileno': fileno,
            'tsp_pos': tsp_pos,
            'tsp_max': tsp_max,
            'ppm': ppm_shifted,
            'magnitude': magnitude
        })
        
        # Plot full spectrum (normalized by TSP max for comparison)
        normalized = magnitude / tsp_max
        ax_full.plot(ppm_shifted, normalized, color=colors[i], 
                    label=f'File {fileno}', linewidth=0.8, alpha=0.8)
        
        # Plot TSP region (normalized)
        mask_tsp_zoom = (ppm_shifted >= -0.2) & (ppm_shifted <= 0.2)
        ax_tsp.plot(ppm_shifted[mask_tsp_zoom], normalized[mask_tsp_zoom], 
                   color=colors[i], linewidth=1.5, alpha=0.8)
        
        # Plot methyl region (normalized)
        mask_methyl = (ppm_shifted >= 1.3) & (ppm_shifted <= 1.6)
        ax_methyl.plot(ppm_shifted[mask_methyl], normalized[mask_methyl],
                      color=colors[i], linewidth=1.5, alpha=0.8)
        
        # Plot CH region (normalized)
        mask_ch = (ppm_shifted >= 3.6) & (ppm_shifted <= 4.0)
        ax_ch.plot(ppm_shifted[mask_ch], normalized[mask_ch],
                  color=colors[i], linewidth=1.5, alpha=0.8)
    
    # Configure full spectrum plot
    ax_full.set_xlabel('Chemical Shift (ppm)', fontsize=12)
    ax_full.set_ylabel('Normalized Intensity (Magnitude/TSP)', fontsize=12)
    ax_full.set_title('Full Spectrum - Alanine Reference Library (TSP-normalized)', 
                     fontsize=14, fontweight='bold')
    ax_full.legend(loc='upper right', fontsize=9, ncol=2)
    ax_full.grid(True, alpha=0.3)
    ax_full.set_xlim(-1, 10)
    ax_full.axhline(y=0, color='k', linestyle='-', alpha=0.3)
    
    # Configure TSP region plot
    ax_tsp.set_xlabel('Chemical Shift (ppm)', fontsize=11)
    ax_tsp.set_ylabel('Normalized Intensity', fontsize=11)
    ax_tsp.set_title('TSP Region', fontsize=12, fontweight='bold')
    ax_tsp.grid(True, alpha=0.3)
    ax_tsp.axvline(x=0, color='k', linestyle='--', alpha=0.5)
    
    # Configure methyl region plot
    ax_methyl.set_xlabel('Chemical Shift (ppm)', fontsize=11)
    ax_methyl.set_ylabel('Normalized Intensity', fontsize=11)
    ax_methyl.set_title('Alanine Methyl (CH₃) Region', fontsize=12, fontweight='bold')
    ax_methyl.grid(True, alpha=0.3)
    ax_methyl.axvline(x=1.47, color='k', linestyle='--', alpha=0.5)
    
    # Configure CH region plot
    ax_ch.set_xlabel('Chemical Shift (ppm)', fontsize=11)
    ax_ch.set_ylabel('Normalized Intensity', fontsize=11)
    ax_ch.set_title('Alanine CH Region', fontsize=12, fontweight='bold')
    ax_ch.grid(True, alpha=0.3)
    ax_ch.axvline(x=3.78, color='k', linestyle='--', alpha=0.5)
    
    # Create summary table
    ax_summary.axis('off')
    
    # Get File 10 reference
    ref_data = next(d for d in tsp_data if d['fileno'] == 10)
    
    # Build table data
    table_text = "Scale(TSP) = TSP_sample / TSP_File10\n\n"
    table_text += f"{'File':<8} {'Scale(TSP)':<12} {'vs Table S5':<12}\n"
    table_text += "-" * 35 + "\n"
    
    table_s5 = {10: 1.0000, 20: 0.9932, 30: 0.6607, 40: 0.8964, 
                50: 0.8992, 60: 0.9546, 70: 0.8852}
    
    for d in tsp_data:
        scale = d['tsp_max'] / ref_data['tsp_max']
        expected = table_s5.get(d['fileno'], 'N/A')
        table_text += f"File {d['fileno']:<4} {scale:<12.4f} {expected:<12}\n"
    
    table_text += "\n" + "=" * 35 + "\n"
    table_text += "Problem: Our Scale(TSP) varies 0.5-1.5\n"
    table_text += "but Table S5 shows only 0.66-1.00\n"
    
    ax_summary.text(0.1, 0.9, table_text, transform=ax_summary.transAxes,
                   fontsize=10, fontfamily='monospace', verticalalignment='top')
    
    plt.tight_layout()
    plt.savefig('alanine_full_range_magnitude.png', dpi=150, bbox_inches='tight')
    print("Saved: alanine_full_range_magnitude.png")
    
    # Print detailed summary
    print("\n" + "="*80)
    print("ALANINE FULL-RANGE SPECTRA - Magnitude Mode")
    print("="*80)
    print(f"{'File':<8} {'TSP Max':<15} {'Scale(TSP)':<12} {'Expected':<12} {'Ratio':<10}")
    print("-"*80)
    
    for d in tsp_data:
        scale = d['tsp_max'] / ref_data['tsp_max']
        expected = table_s5.get(d['fileno'], None)
        ratio = scale/expected if expected else None
        expected_str = f"{expected:.4f}" if expected else "N/A"
        ratio_str = f"{ratio:.2f}" if ratio else "N/A"
        print(f"File {d['fileno']:<4} {d['tsp_max']:<15.3e} {scale:<12.4f} "
              f"{expected_str:<12} {ratio_str:<10}")
    
    print("="*80)
    
    plt.close()


if __name__ == "__main__":
    plot_full_range_spectra()