plot_model_mixtures.py

#!/usr/bin/env python3
"""
Plot all model mixture spectra to help identify M1-M6
"""

import sys
sys.path.insert(0, '.')
import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
from quantify_all_metabolites_v3_ref20 import read_and_process, find_tsp_peak

def plot_all_mixtures():
    folder_path = "raw_data/Model_Mixtures-M1-M6_JCAMP-DX"
    
    # Get all .dx files and sort numerically
    files = sorted([f for f in os.listdir(folder_path) if f.endswith('.dx')],
                   key=lambda x: int(x.replace('.dx', '')))
    
    print(f"Found {len(files)} files: {files}")
    
    # Create a large figure with subplots
    n_files = len(files)
    n_cols = 6
    n_rows = (n_files + n_cols - 1) // n_cols
    
    fig = plt.figure(figsize=(24, 4 * n_rows))
    
    for idx, filename in enumerate(files):
        ax = fig.add_subplot(n_rows, n_cols, idx + 1)
        
        try:
            # Read and process
            ppm, spec = read_and_process(os.path.join(folder_path, filename))
            tsp = find_tsp_peak(ppm, spec)
            ppm_corr = ppm - tsp
            
            # Plot full spectrum (0.5-5.5 ppm)
            mask = (ppm_corr >= 0.5) & (ppm_corr <= 5.5)
            ax.plot(ppm_corr[mask], spec[mask], linewidth=0.8)
            
            # Styling
            ax.set_xlim(5.5, 0.5)
            ax.set_title(f'{filename.replace(".dx", "")}', fontsize=12, fontweight='bold')
            ax.set_xlabel('ppm', fontsize=9)
            ax.set_ylabel('Intensity', fontsize=9)
            ax.tick_params(labelsize=8)
            ax.grid(True, alpha=0.2)
            
        except Exception as e:
            ax.text(0.5, 0.5, f'Error: {e}', ha='center', va='center', transform=ax.transAxes)
            ax.set_title(f'{filename}', fontsize=10)
    
    plt.suptitle('Model Mixtures M1-M6 - All Spectra (TSP-corrected)', 
                 fontsize=16, fontweight='bold', y=1.00)
    plt.tight_layout()
    
    output_path = 'model_mixtures_all_spectra.png'
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"\nSaved to: {output_path}")
    plt.close()
    
    # Also create individual detailed plots for each file
    print("\nGenerating individual detailed plots...")
    os.makedirs('model_mixtures_individual', exist_ok=True)
    
    for filename in files:
        fig, axes = plt.subplots(2, 2, figsize=(16, 10))
        
        try:
            ppm, spec = read_and_process(os.path.join(folder_path, filename))
            tsp = find_tsp_peak(ppm, spec)
            ppm_corr = ppm - tsp
            
            # Full spectrum
            ax = axes[0, 0]
            mask = (ppm_corr >= 0.5) & (ppm_corr <= 5.5)
            ax.plot(ppm_corr[mask], spec[mask], linewidth=1)
            ax.set_xlim(5.5, 0.5)
            ax.set_title('Full Spectrum (0.5-5.5 ppm)')
            ax.set_xlabel('Chemical Shift (ppm)')
            ax.set_ylabel('Intensity')
            ax.xaxis.set_minor_locator(MultipleLocator(0.1))
            ax.grid(True, alpha=0.3)
            
            # Aliphatic region (0.5-3.0 ppm)
            ax = axes[0, 1]
            mask = (ppm_corr >= 0.5) & (ppm_corr <= 3.0)
            ax.plot(ppm_corr[mask], spec[mask], linewidth=1, color='blue')
            ax.set_xlim(3.0, 0.5)
            ax.set_title('Aliphatic Region (0.5-3.0 ppm)')
            ax.set_xlabel('Chemical Shift (ppm)')
            ax.set_ylabel('Intensity')
            ax.xaxis.set_minor_locator(MultipleLocator(0.05))
            ax.grid(True, alpha=0.3)
            
            # CH/OH region (3.0-4.5 ppm)
            ax = axes[1, 0]
            mask = (ppm_corr >= 3.0) & (ppm_corr <= 4.5)
            ax.plot(ppm_corr[mask], spec[mask], linewidth=1, color='green')
            ax.set_xlim(4.5, 3.0)
            ax.set_title('CH/OH Region (3.0-4.5 ppm)')
            ax.set_xlabel('Chemical Shift (ppm)')
            ax.set_ylabel('Intensity')
            ax.xaxis.set_minor_locator(MultipleLocator(0.05))
            ax.grid(True, alpha=0.3)
            
            # Aromatic region (6.5-8.0 ppm)
            ax = axes[1, 1]
            mask = (ppm_corr >= 6.5) & (ppm_corr <= 8.0)
            if np.sum(mask) > 0:
                ax.plot(ppm_corr[mask], spec[mask], linewidth=1, color='red')
                ax.set_xlim(8.0, 6.5)
                ax.set_title('Aromatic Region (6.5-8.0 ppm)')
            else:
                ax.text(0.5, 0.5, 'No data in aromatic region', 
                       ha='center', va='center', transform=ax.transAxes)
                ax.set_title('Aromatic Region - No Data')
            ax.set_xlabel('Chemical Shift (ppm)')
            ax.set_ylabel('Intensity')
            ax.grid(True, alpha=0.3)
            
            plt.suptitle(f'Model Mixture - File {filename.replace(".dx", "")}\nTSP Correction: +{-tsp:.4f} ppm', 
                        fontsize=14, fontweight='bold')
            plt.tight_layout()
            
            output_file = f'model_mixtures_individual/{filename.replace(".dx", "")}_detailed.png'
            plt.savefig(output_file, dpi=150, bbox_inches='tight')
            plt.close()
            
        except Exception as e:
            print(f"Error processing {filename}: {e}")
            plt.close()
    
    print(f"Individual plots saved to: model_mixtures_individual/")

if __name__ == "__main__":
    plot_all_mixtures()