plot_alanine_full_range_correct.py

#!/usr/bin/env python3
"""
Plot full-range magnitude spectra for alanine files
Using CORRECT ppm scale ($O1/$SWH method)
"""

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


def process_spectrum(dic, data):
    """Process to magnitude spectrum"""
    real = data[0]
    imag = data[1] if len(data) > 1 else np.zeros_like(real)
    return np.sqrt(real**2 + imag**2)


def get_ppm_scale(dic, n_points):
    """Calculate ppm scale from $O1 and $SWH (correct method)"""
    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_correct():
    """Plot full-range spectra using correct ppm scale"""
    
    base_dir = "raw_data/Reference_Raw_Date_JCAMP-DX/Alanine-Reference"
    files = [10, 20, 30, 40, 50, 60, 70]
    
    fig = plt.figure(figsize=(16, 12))
    gs = fig.add_gridspec(3, 2, height_ratios=[2, 1, 1])
    
    ax_full = fig.add_subplot(gs[0, :])
    ax_tsp = fig.add_subplot(gs[1, 0])
    ax_methyl = fig.add_subplot(gs[1, 1])
    ax_ch = fig.add_subplot(gs[2, 0])
    ax_summary = fig.add_subplot(gs[2, 1])
    
    colors = plt.cm.tab10(np.linspace(0, 1, len(files)))
    
    tsp_data = []
    
    for i, fileno in enumerate(files):
        filepath = os.path.join(base_dir, f"{fileno}.dx")
        if not os.path.exists(filepath):
            continue
        
        dic, data = ng.jcampdx.read(filepath)
        magnitude = process_spectrum(dic, data)
        ppm = get_ppm_scale(dic, len(magnitude))
        
        # Find TSP position
        mask_tsp = (ppm >= -0.5) & (ppm <= 0.5)
        tsp_idx = np.argmax(magnitude[mask_tsp])
        tsp_pos = ppm[mask_tsp][tsp_idx]
        ppm_shifted = ppm - tsp_pos
        
        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 normalized by TSP
        normalized = magnitude / tsp_max
        ax_full.plot(ppm_shifted, normalized, color=colors[i], 
                    label=f'File {fileno}', linewidth=0.8, alpha=0.8)
        
        # TSP region zoom
        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)
        
        # Methyl region
        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)
        
        # CH region
        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 plots
    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 (CORRECT ppm scale)', 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)
    
    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)
    
    ax_methyl.set_xlabel('Chemical Shift (ppm)', fontsize=11)
    ax_methyl.set_ylabel('Normalized Intensity', fontsize=11)
    ax_methyl.set_title('Alanine 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)
    
    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)
    
    # Summary table
    ax_summary.axis('off')
    
    ref_data = next(d for d in tsp_data if d['fileno'] == 10)
    
    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"
    
    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_CORRECT.png', dpi=150, bbox_inches='tight')
    print("Saved: alanine_full_range_CORRECT.png")
    
    # Print summary
    print("\n" + "="*80)
    print("ALANINE FULL-RANGE SPECTRA - CORRECT ppm scale")
    print("="*80)
    print(f"{'File':<8} {'TSP Max':<15} {'Scale(TSP)':<12} {'Expected':<12}")
    print("-"*80)
    
    for d in tsp_data:
        scale = d['tsp_max'] / ref_data['tsp_max']
        expected = table_s5.get(d['fileno'], None)
        expected_str = f"{expected:.4f}" if expected else "N/A"
        print(f"File {d['fileno']:<4} {d['tsp_max']:<15.3e} {scale:<12.4f} {expected_str:<12}")
    
    print("="*80)
    plt.close()


if __name__ == "__main__":
    plot_full_range_correct()