plot_tsp_region_only.py

#!/usr/bin/env python3
"""
Plot only the TSP region for all alanine files
Using correct ppm scale - NO normalization
"""

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_tsp_region():
    """Plot TSP region only - no normalization"""
    
    base_dir = "raw_data/Reference_Raw_Date_JCAMP-DX/Alanine-Reference"
    files = [10, 20, 30, 40, 50, 60, 70]
    
    fig, axes = plt.subplots(2, 4, figsize=(18, 10))
    axes = axes.flatten()
    
    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
        
        # Extract TSP region
        mask_region = (ppm_shifted >= -0.3) & (ppm_shifted <= 0.3)
        ppm_region = ppm_shifted[mask_region]
        mag_region = magnitude[mask_region]
        
        tsp_max = np.max(mag_region)
        tsp_data.append({
            'fileno': fileno,
            'tsp_max': tsp_max,
            'ppm': ppm_region,
            'signal': mag_region
        })
        
        # Individual plot
        ax = axes[i]
        ax.plot(ppm_region, mag_region, color=colors[i], linewidth=1.5)
        ax.set_xlabel('Chemical Shift (ppm)', fontsize=10)
        ax.set_ylabel('Intensity (Magnitude)', fontsize=10)
        ax.set_title(f'File {fileno} - TSP Signal\nMax: {tsp_max:.3e}', fontsize=11, fontweight='bold')
        ax.grid(True, alpha=0.3)
        ax.axvline(x=0, color='k', linestyle='--', alpha=0.5)
    
    # Use last subplot for comparison
    ax_compare = axes[-1]
    for i, d in enumerate(tsp_data):
        ax_compare.plot(d['ppm'], d['signal'], color=colors[i], 
                       label=f'File {d["fileno"]}', linewidth=1.5, alpha=0.8)
    ax_compare.set_xlabel('Chemical Shift (ppm)', fontsize=10)
    ax_compare.set_ylabel('Intensity (Magnitude)', fontsize=10)
    ax_compare.set_title('All Files - TSP Comparison\n(No normalization)', fontsize=11, fontweight='bold')
    ax_compare.legend(loc='upper right', fontsize=8)
    ax_compare.grid(True, alpha=0.3)
    ax_compare.axvline(x=0, color='k', linestyle='--', alpha=0.5)
    
    plt.tight_layout()
    plt.savefig('tsp_region_individual.png', dpi=150, bbox_inches='tight')
    print("Saved: tsp_region_individual.png")
    
    # Print summary table
    print("\n" + "="*70)
    print("TSP SIGNAL SUMMARY - Raw Magnitude Values")
    print("="*70)
    print(f"{'File':<8} {'TSP Max':<18} {'Ratio vs File 10':<15}")
    print("-"*70)
    
    ref_tsp = tsp_data[0]['tsp_max']  # File 10
    
    for d in tsp_data:
        ratio = d['tsp_max'] / ref_tsp
        print(f"File {d['fileno']:<4} {d['tsp_max']:<18.3e} {ratio:<15.2f}")
    
    print("="*70)
    print(f"\nFile 10 (reference): {ref_tsp:.3e}")
    print(f"File 40 TSP max: {tsp_data[3]['tsp_max']:.3e}")
    print(f"File 70 TSP max: {tsp_data[6]['tsp_max']:.3e}")
    print(f"\nRange: {min([d['tsp_max'] for d in tsp_data])/ref_tsp:.1f}x to {max([d['tsp_max'] for d in tsp_data])/ref_tsp:.1f}x File 10")
    
    plt.close()


if __name__ == "__main__":
    plot_tsp_region()