PtolemyGUI/Cleopatra/PlotWindow.py

#!/usr/bin/python3

from PyQt6.QtWidgets import (
  QGridLayout, QWidget, QCheckBox
)

import numpy as np

import matplotlib.pyplot as plt
from matplotlib.backends.backend_qtagg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas

# Set backend to a Qt-compatible one
plt.switch_backend('QtAgg')  # Or use 'Qt5Agg' if there are still issues

class FitPlotWindow(QWidget):
  def __init__(self, windowTitle):
    super().__init__()

    self.setWindowTitle(windowTitle)
    self.resize(800, 600)

    self.default_colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

    self.log_scale_checkbox = QCheckBox("Use Log Scale for Y-Axis")
    self.log_scale_checkbox.setChecked(True)
    self.log_scale_checkbox.stateChanged.connect(self.plot_Fit)

    self.gridline_checkbox = QCheckBox("Show Gridlines")
    self.gridline_checkbox.stateChanged.connect(self.plot_Fit)

    self.figure, self.ax = plt.subplots()
    self.canvas = FigureCanvas(self.figure)
    self.toolbar = NavigationToolbar(self.canvas, self)

    layout = QGridLayout(self)
    layout.addWidget(self.toolbar, 0, 0, 1, 3)
    layout.addWidget(self.log_scale_checkbox, 1, 0)
    layout.addWidget(self.gridline_checkbox, 1, 1)
    layout.addWidget(self.canvas, 2, 0, 5, 3)

    self.setLayout(layout)

  def set_data(self, ID, expData, fitOption, dataName_list, xData, yData_list, para, perr, chi_square ):
    self.x_exp = expData[ID][:, 0]
    self.x_err = expData[ID][:, 1]
    self.y_exp = expData[ID][:, 2]
    self.y_err = expData[ID][:, 3]
    self.dataName = dataName_list[ID]
    self.fitOption = fitOption[ID]

    self.xData = xData
    self.yData_list = yData_list

    self.para = para
    self.perr = perr
    self.chi_square = chi_square

  def plot_Fit(self):
    self.ax.clear()

    self.ax.errorbar(self.x_exp, self.y_exp, xerr=self.x_err, yerr=self.y_err, 
                    fmt='x', label='Experimental Data', color='black', markersize = 15, elinewidth=2)

    fitTheory = []
    fitTheory_lower = []
    fitTheory_upper = []
    
    for k in range(len(self.fitOption)):

      xsecIDStr = self.fitOption[k].strip()
      xsecID = [int(part.strip()) for part in xsecIDStr.split('+')] if '+' in xsecIDStr else [int(xsecIDStr)]

      fitTheory.append(np.zeros_like(self.xData))
      fitTheory_upper.append(np.zeros_like(self.xData))
      fitTheory_lower.append(np.zeros_like(self.xData))
      
      for id in xsecID:
        fitTheory[k] += self.para[k] * np.interp(self.xData, self.xData, self.yData_list[id])
        fitTheory_upper[k] += (self.para[k] + self.perr[k]) * np.interp(self.xData, self.xData, self.yData_list[id])
        fitTheory_lower[k] += (self.para[k] - self.perr[k]) * np.interp(self.xData, self.xData, self.yData_list[id])


    for i, fit in enumerate(fitTheory):
      self.ax.plot(self.xData, fit, label=f'Chi2:{self.chi_square[i]:.3f} | Xsec:{self.fitOption[i]}')
      self.ax.fill_between(self.xData, fitTheory_lower[i], fitTheory_upper[i], alpha=0.2)

      self.ax.text(0.05, 0.1 + 0.05*i, f"Xsec-{self.fitOption[i].strip()}: {', '.join([f'{x:.3f}' for x in self.para[i]])} +/- {', '.join([f'{x:.3f}' for x in self.perr[i]])}" ,
                transform=plt.gca().transAxes, fontsize=12, 
                verticalalignment='bottom', horizontalalignment='left', color=self.default_colors[i])

    # Replace plt.title() with plt.text() to position the title inside the plot
    self.ax.text(0.05, 0.05, f'Fit for Exp Data : {self.dataName}', transform=plt.gca().transAxes,
        fontsize=12, verticalalignment='bottom', horizontalalignment='left', color='black')


    # Plot decorator
    # Apply log scale for y-axis if selected
    if self.log_scale_checkbox.isChecked():
      self.ax.set_yscale('log')
    else:
      self.ax.set_yscale('linear')


    if self.gridline_checkbox.isChecked():
      self.ax.grid(True, which='both', linestyle='--', linewidth=0.5, color='gray')
    else:
      self.ax.grid(False)


    self.ax.set_xlabel(r'$\theta_{cm}$ [deg]')
    self.ax.set_ylabel(r'd$\sigma$/d$\Omega$ [deg]')
    self.ax.legend(loc='upper right', frameon=True)

    self.ax.autoscale(enable=True, axis='x', tight=True)
    self.figure.subplots_adjust(left=0.1, right=0.95, top=0.95, bottom=0.1)

    self.canvas.draw_idle()