ploty_offsets_align_plots.py

# Download the required Data
#!wget https://ucmerced.box.com/shared/static/1qivuf3teov6qm9oealjtrjgw1glpbwa.gz \ -O sample_data/f28_data.tar.gz && mkdir sample_data/f28_data && \
# tar -zxvf sample_data/f28_data.tar.gz -C sample_data/f28_data

import numpy as np
import plotly.io as pio
import plotly.graph_objs as go
from plotly.subplots import make_subplots

# Changing Default Template to Plotly White
pio.templates.default = "plotly_white"

# Constants
FWHM_0K = np.zeros(4)
FWHM_300K = np.zeros(4)
label = ["Gas", "PCM", "Strip", "Explicit"]  # DMSO

# Conversion Factors
INVCM_TO_EV = float(0.00012398)

# File Paths
PATHS = [
    "sample_data/f28_data/nbdnh2_gas_frame28_vibronic_ems.dat",
    "sample_data/f28_data/nbdnh2_pcm_frame28_vibronic_ems.dat",
    "sample_data/f28_data/nbdnh2_stripped_dmso_frame28_vibronic_ems.dat",  # exAIMD
    "sample_data/f28_data/nbdnh2_frame28_fc_qm47_traj2_run3.dat",  # exAIMD
]


# Important Functions
# Load and process data
def load_process_data(paths):
    """
    The `load_process_data` function is responsible for loading and
    processing the data from the specified file paths.
    """
    data = [np.loadtxt(path, skiprows=7, usecols=(0, 1, 2)) for path in paths]
    curves_0K = [(d[:, 0] * INVCM_TO_EV, d[:, 1] / d[:, 1].max()) for d in data]
    curves_300K = [(d[:, 0] * INVCM_TO_EV, d[:, 2] / d[:, 2].max()) for d in data]
    return curves_0K, curves_300K


# FWHM Function
def FWHM(x: np.ndarray, y: np.ndarray, height: float = 0.5) -> float:
    """
    The `FWHM` function is calculating the full width at half maximum
    (FWHM) of a given curve represented by `x` and `y` arrays.
    The function takes in three arguments: `x` and `y` arrays
    representing the curve, and `height` which is the fraction of the
    maximum height at which the FWHM is calculated (default value is 0.5).
    """

    height_half_max = np.max(y) * height
    index_max = np.argmax(y)
    x_low = np.interp(height_half_max, y[: index_max + 1], x[: index_max + 1])
    x_high = np.interp(height_half_max, np.flip(y[index_max:]), np.flip(x[index_max:]))

    return abs(x_high - x_low)


# Align curves function
def align_curves(curves):
    """
    The `align_curves` function is used to align the x-axis of the curves.
    It takes in a list of curves, where each curve is represented by a tuple of x and y arrays.
    """
    x_ref = curves[0][0][0]
    return [(c[0] - c[0][0] + x_ref, c[1]) for c in curves]


# Load Data And Align Curves
curves_0K, curves_300K = load_process_data(PATHS)
curves_0K = align_curves(curves_0K)
curves_300K = align_curves(curves_300K)

# Compute FWHM of All
for idx, curve_0K in enumerate(curves_0K):
    FWHM_0K[idx] = FWHM(curve_0K[0], curve_0K[1])

for idx, curve_300K in enumerate(curves_300K):
    FWHM_300K[idx] = FWHM(curve_300K[0], curve_300K[1])

# Create subplots
fig = make_subplots(
    rows=1, cols=2, subplot_titles=("Temperature: 0K", "Temperature: 300K")
)

# Define color map
color_map = {"Gas": "#A084E8", "PCM": "orange", "Strip": "red", "Explicit": "black"}

# Add traces for 0K temperature
for idx, curve in enumerate(curves_0K):
    fig.add_trace(
        go.Scatter(
            x=curve[0],
            y=curve[1],
            mode="lines",
            name=f"{label[idx]}: {FWHM_0K[idx]:.3f}",
            line=dict(color=color_map[label[idx]], dash="dash"),
        ),
        row=1,
        col=1,
    )

# Add traces for 300K temperature
for idx, curve in enumerate(curves_300K):
    fig.add_trace(
        go.Scatter(
            x=curve[0],
            y=curve[1],
            mode="lines",
            name=f"{label[idx]}: {FWHM_300K[idx]:.3f}",
            line=dict(color=color_map[label[idx]]),
        ),
        row=1,
        col=2,
    )

# Update xaxis titles and range
fig.update_xaxes(title_text="Energy (eV)", range=[2.6, 3.8], nticks=6, row=1, col=1)
fig.update_xaxes(title_text="Energy (eV)", range=[2.6, 3.8], nticks=6, row=1, col=2)

# Update yaxis titles and range
fig.update_yaxes(title_text="Intensity", range=[0.0, 1.1], nticks=4, row=1, col=1)
fig.update_yaxes(range=[0.0, 1.1], nticks=3, row=1, col=2)

# Add layout
# fig.update_layout(height=600, width=1200,
#                  title_text="Energy vs Intensity")
fig.to_image("allign_offset_plots.png", scale=3)
fig.show()