nmreval/src/gui_qt/dsc/glass_dialog.py

from itertools import cycle

from numpy import array
from pyqtgraph import mkPen, mkBrush

from nmreval.lib.colors import Tab10
from ..Qt import QtWidgets, QtCore
from .._py.tnmh_dialog import Ui_Dialog
from ..lib.pg_objects import PlotItem, RegionItem

from nmreval.data import DSC


class TgCalculator(QtWidgets.QDialog, Ui_Dialog):
    def __init__(self, management, parent=None):
        super().__init__(parent=parent)

        self.setupUi(self)
        self._management = management
        self._colors = cycle(Tab10)

        self._dsc = {}
        self._tg_baselines = {}
        self._tg_value_plot = {}
        self.limits = RegionItem(), RegionItem()
        for lim in self.limits:
            self.graphicsView.addItem(lim)

        self.add_sets()

        self.listWidget.itemClicked.connect(self.show_tg_values)

    def __call__(self):
        self.clear()
        self.add_sets()
        # s.valuesChanged.connect(self.management.shift_scale)

    def clear(self):
        self.listWidget.clear()
        for plot, data in self._dsc.values():
            self.graphicsView.removeItem(plot)
        self._dsc = {}

        for v in self._tg_baselines.values():
            self.graphicsView.removeItem(v)
        self._tg_baselines = {}

        for v in self._tg_value_plot:
            self.graphicsView.removeItem(v)
        self._tg_value_plot = []

    def add_sets(self):
        min_x = 10_000_000
        max_x = -10_000_000
        for (key, name), c in zip(self._management.active_sets, self._colors):
            data = self._management[key].data
            if not isinstance(data, DSC):
                continue

            min_x = min(min_x, data.x.min())
            max_x = max(max_x, data.x.max())

            item = QtWidgets.QListWidgetItem(name)
            item.setData(QtCore.Qt.UserRole, key)
            item.setForeground(mkBrush(c.rgb()))
            self.listWidget.addItem(item)

            plot = PlotItem(x=data.x, y=data.y, pen=mkPen(c.rgb()))
            self.graphicsView.addItem(plot)
            self._dsc[key] = (plot, data)

        dist = max_x - min_x
        self.limits[0].setRegion((min_x, min_x+0.1*dist))
        self.limits[1].setRegion((max_x-0.1*dist, max_x))

    @QtCore.pyqtSlot(name='on_pushButton_clicked')
    def calc_tg(self):
        for v in self._tg_baselines.values():
            for line in v:
                self.graphicsView.removeItem(line)
        self._tg_baselines = {}

        for v in self._tg_value_plot:
            self.graphicsView.removeItem(v)

        baselines = tuple(lim.getRegion() for lim in self.limits)
        if baselines[0][0] > baselines[1][0]:
            baselines = baselines[1], baselines[0]

        for idx in range(self.listWidget.count()):
            item = self.listWidget.item(idx)
            key = item.data(QtCore.Qt.UserRole)
            color = item.foreground()
            _, data = self._dsc[key]
            tg_results, glass, liquid, tangent = data.glass_transition(*baselines)

            lines = []
            for line in (glass, liquid, tangent):
                plt = PlotItem(line[:, 0], line[:, 1])
                plt.set_line(style=2, color=color.color())
                self.graphicsView.addItem(plt)
                lines.append(plt)
            self._tg_baselines[key] = lines
            tg_plot = PlotItem(array(list(tg_results.values())), pen=None, symbolBrush=color, symbol='o')
            self._tg_value_plot[key] = (tg_results, tg_plot)
            self.graphicsView.addItem(tg_plot)

    def show_tg_values(self, item):
        values, _ = self._tg_value_plot.get(item.data(QtCore.Qt.UserRole), ({}, None))

        if values is not None:
            self.listWidget_2.clear()
            for name, pos in values.items():
                tg_item = QtWidgets.QListWidgetItem(f'{name.capitalize()}: {pos[0]:.2f} K')
                self.listWidget_2.addItem(tg_item)

    @QtCore.pyqtSlot(name='on_pushButton_2_clicked')
    def get_fictive(self):
        baselines = tuple(lim.getRegion() for lim in self.limits)
        if baselines[0][0] > baselines[1][0]:
            baselines = baselines[1], baselines[0]

        for idx in range(self.listWidget.count()):
            item = self.listWidget.item(idx)
            key = item.data(QtCore.Qt.UserRole)
            color = item.foreground()
            _, data = self._dsc[key]
            cp, tg = data.get_fictive_cp(*baselines)

            plt = PlotItem(cp.x, cp.y)
            plt.set_line(color=color.color())
            self.graphicsView_2.addItem(plt)

            self._tg_value_plot[key][0]['fictive'] = (tg, 0)
            tg_plot = PlotItem(array(list(self._tg_value_plot[key][0].values())), pen=None, symbolBrush=color, symbol='o')
            self.graphicsView.addItem(tg_plot)

            res = data.calculate_tnmh([100, 0.5, 1, 2e5], *baselines, return_fictive=False)
            plt = PlotItem(res.x, res.y)
            print(res.pprint())
            plt.set_line(style=2, color=color.color())
            self.graphicsView_2.addItem(plt)

    def close(self) -> bool:
        self.clear()
        return super().close()