diff_plotting_chunks.py

# -*- coding: utf-8 -*-
"""
This script calculates and plots the Mean Squared Displacement (MSD) for different protein simulations
from various force fields. It reads data from multiple files, processes it, and generates a 3x2 subplot
grid to compare the MSD over time for each force field. The output is a high-resolution figure
suitable for academic publications.
"""

import os
import numpy as np
import matplotlib.pyplot as plt

# --- Plotting Style Configuration ---
# Set global font sizes for better readability in publications.
plt.rcParams['font.size'] = 14
plt.rcParams['axes.labelsize'] = 16
plt.rcParams['xtick.labelsize'] = 14
plt.rcParams['ytick.labelsize'] = 14
# --- End of Configuration ---

# List of folders containing the simulation data for different force fields.
folder_paths = [
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/20_AMBER/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/12_AA/1_monomer/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/13_SIRAH/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/9_Martini3/3_303K/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/10_Martini2_nonPW/3_303K/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/11_Martini2_PW/3_303K/"
]

# Specific force field names for plot titles as requested
ff_names = ['AMBER', 'CHARMM', 'SIRAH hybrid', 'M3', 'M2', 'M2PW']


def parse_file(filename):
    """
    Parses a single data file to extract time (tau) and MSD values.

    Args:
        filename (str): The path to the input file.

    Returns:
        tuple: A tuple containing two lists, tau (time) and msd (Mean Squared Displacement).
    """
    tau = []
    msd = []
    with open(filename, 'r') as file:
        lines = file.readlines()
        for line in lines:
            if line.startswith("#") or line.startswith("@"):
                continue
            data = line.split()
            if len(data) >= 2:
                tau.append(float(data[0]))
                msd.append(float(data[1]))
    return tau, msd


# Create a 2x3 subplot grid with a more square aspect ratio for each plot
fig, axes = plt.subplots(2, 3, figsize=(18, 10), sharex=False, sharey=False)
axes = axes.flatten()  # Flatten the 2D array of axes to easily iterate over it

# Loop through each folder path and its corresponding subplot
for i, folder_path in enumerate(folder_paths):
    ax = axes[i]
    time = None
    all_msd_data = []

    # This block attempts to find data files. It will only work if the script
    # is run in an environment where the paths are accessible.
    try:
        # Check for "Take_" subdirectories
        takes = [os.path.join(folder_path, d) for d in os.listdir(folder_path) if
                 os.path.isdir(os.path.join(folder_path, d)) and 'Take_' in d]

        if not takes:  # If no "Take_" folders, assume data is directly in the folder_path
            takes = [folder_path]

        for take_path in takes:
            # Generate the list of filenames
            filenames = []
            if os.path.exists(take_path):
                for f in os.listdir(take_path):
                    if f.startswith('diff_RbsB_SpeA_FkpA_') and f.endswith('.xvg'):
                        filenames.append(os.path.join(take_path, f))

            for filename in filenames:
                if os.path.exists(filename):
                    tau, msd = parse_file(filename)
                    if time is None and tau:
                        time = np.array(tau) / 1000  # Convert ns to microseconds
                    if tau and len(time) == len(msd):
                        all_msd_data.append(msd)
                    else:
                        if tau:  # Only print warning if tau has data
                            print(f"Warning: Length mismatch in file {filename} (time: {len(time)}, msd: {len(msd)})")
                else:
                    print(f"Warning: File not found {filename}")
    except FileNotFoundError:
        print(f"Directory not found: {folder_path}. This subplot will be empty.")

    # Plotting the data for the current force field
    if time is not None and all_msd_data:
        for msd in all_msd_data:
            ax.plot(time, msd)

    # Highlight the region from 0.02 to 0.04 microseconds
    ax.axvspan(0.02, 0.04, color='gray', alpha=0.3)

    # Set title for the subplot with larger, bold font
    ax.set_title(ff_names[i], fontsize=18, fontweight='bold')

    # Set common labels (size is now controlled by rcParams)
    ax.set_xlabel(r'$\tau$ ($\mu$s)')
    ax.set_ylabel('MSD (nm$^2$)')
    ax.set_xlim(0, 0.1)
    ax.set_ylim(0, 0.09)

# Overall title for the entire figure with larger, bold font
# fig.suptitle('MSD Comparison Across Different Force Fields', fontsize=22, fontweight='bold')

# Adjust layout to prevent titles and labels from overlapping
plt.tight_layout(rect=[0, 0.03, 1, 0.95])

# Save the plot to a file with high resolution and the new filename
plt.savefig('msd_subplots_lines_all_reps_andFFs.png', dpi=600)

# Show the plot
plt.show()

print("Plot saved as msd_subplots_lines_all_reps_andFFs.png")