Diffusion_Water_plot.py

import pandas as pd
import matplotlib.pyplot as plt

# --- Function Definition ---

def plot_diffusion_coefficients(csv_path):
    """
    Reads diffusion coefficient data from a CSV, converts units,
    prints summary statistics, and generates the original-style plot.

    Args:
        csv_path (str): The file path for the input CSV data.
                        The CSV should have force fields as columns.
    """
    # --- 1. Data Loading and Processing ---
    try:
        # Read the CSV file, skipping the first (index) column.
        data = pd.read_csv(csv_path).iloc[:, 1:]
    except FileNotFoundError:
        print(f"Error: The file '{csv_path}' was not found.")
        print("Please ensure the script is in the same directory as your data file.")
        return

    # The average and standard deviation for each force field were calculated.
    means = data.mean()
    stds = data.std()

    # --- 2. Unit Conversion ---

    # The conversion factor was defined to change units from 1e-5 cm^2/s
    # to 1e-7 nm^2/s. This requires multiplying by 100.
    conversion_factor = 100

    # The conversion factor was applied to the calculated means and standard deviations.
    means_converted = means * conversion_factor
    stds_converted = stds * conversion_factor

    # --- 3. Print Summary Statistics ---
    # The calculated mean and standard deviation for each force field were printed.
    print("--- Diffusion Coefficients (D ± SD) ---")
    for ff_name in means_converted.index:
        print(f"{ff_name}: {means_converted[ff_name]:.2f} ± {stds_converted[ff_name]:.2f} (1 x 10^7 nm^2/s)")
    print("-" * 39)


    # --- 4. Plotting ---

    # A color mapping was created to associate each force field with its specific color.
    color_map = {
        'AMBER': 'olive', 'CHARMM': 'cyan', 'SIRAH hybrid': 'brown',
        'M3': 'magenta', 'M2': 'orange', 'M2PW': 'blue'
    }

    # The desired order for the legend was defined.
    legend_order = ['CHARMM', 'M2PW', 'SIRAH hybrid', 'AMBER', 'M3', 'M2']

    # The figure size for the plot was set.
    plt.figure(figsize=(4.5, 5))

    plt.axhspan(230, 292, color='lime', alpha=0.3, label='Experimental')


    # A loop was used to plot each force field in the specified order.
    # This ensures the legend is generated in the correct sequence.
    for ff_name in legend_order:
        plt.errorbar(
            x=1,
            y=means_converted[ff_name],
            yerr=stds_converted[ff_name],
            fmt='o',  # 'o' specifies a circular marker
            color=color_map[ff_name],
            capsize=5,
            markersize=8,
            label=ff_name
        )

    # The title and axis labels were added. The y-axis label reflects the unit conversion.
    plt.title('Water')
    plt.ylabel('D ($1 \\times 10^7$ nm$^2$/s)')

    # The x-axis tick was set to a single point, as in the original plot.
    plt.xticks([1], ['Force Fields'])

    # A legend was added to the right of the plot area for clear identification.
    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))

    # The layout was adjusted to ensure the legend fits within the saved figure.
    plt.tight_layout(rect=[0, 0, 0.85, 1])

    # The plot was saved as a high-resolution PNG file.
    output_filename = 'diffusion_coefficients_water_plot.png'
    plt.savefig(output_filename, dpi=600, bbox_inches='tight')
    print(f"\nPlot saved successfully as '{output_filename}'")

    # The plot was displayed on the screen.
    plt.show()


# --- Script Execution ---
if __name__ == "__main__":
    # Define the name of the data file.
    data_file = "Diffusion_Water_table.csv"
    # Call the main function to perform the analysis and plotting.
    plot_diffusion_coefficients(data_file)