neuropercolation/evaluation/plot.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Aug 21 16:12:53 2023

@author: timofej
"""

import sys
import os

import pyqtgraph as qt
import pyqtgraph.exporters
from PyQt5.QtGui import QFont


import math as m
import numpy as np

def new_folder(path):
    if not os.path.exists(path):
        os.makedirs(path)
    return path

def __get_color(factor, gamma=0.8):
    frequency=380+factor*(750-380)
    
    lightfrequency = 0.4*(3*np.log10(frequency/2)-2)/4
    
    wavelength = 300/lightfrequency
    
    '''steps of 10Hz: 22 means 220Hz'''

    '''This converts a given wavelength of light to an 
    approximate RGB color value. The wavelength must be given
    in values between 0 and 1 for 0=380nm through 1=750nm
    (789 THz through 400 THz).

    Based on code by Dan Bruton
    http://www.physics.sfasu.edu/astro/color/spectra.html
    '''

    wavelength = float(wavelength)
    if wavelength >= 380 and wavelength <= 440:
        attenuation = 0.3 + 0.7 * (wavelength - 380) / (440 - 380)
        R = ((-(wavelength - 440) / (440 - 380)) * attenuation) ** gamma
        G = 0.0
        B = (1.0 * attenuation) ** gamma
    elif wavelength >= 440 and wavelength <= 490:
        R = 0.0
        G = ((wavelength - 440) / (490 - 440)) ** gamma
        B = 1.0
    elif wavelength >= 490 and wavelength <= 510:
        R = 0.0
        G = 1.0
        B = (-(wavelength - 510) / (510 - 490)) ** gamma
    elif wavelength >= 510 and wavelength <= 580:
        R = ((wavelength - 510) / (580 - 510)) ** gamma
        G = 1.0
        B = 0.0
    elif wavelength >= 580 and wavelength <= 645:
        R = 1.0
        G = (-(wavelength - 645) / (645 - 580)) ** gamma
        B = 0.0
    elif wavelength >= 645 and wavelength <= 750:
        attenuation = 0.3 + 0.7 * (750 - wavelength) / (750 - 645)
        R = (1.0 * attenuation) ** gamma
        G = 0.0
        B = 0.0
    else:
        R = 0.0
        G = 0.0
        B = 0.0
    R *= 255
    G *= 255
    B *= 255
    return (int(R), int(G), int(B))

def plot_execute():
    if sys.flags.interactive != 1 or not hasattr(qt.QtCore, 'PYQT_VERSION'):
        qt.QtGui.QApplication.exec_()

def qtplot(titletext, spaces, vals, names=None, 
           x_tag=f'noise {chr(949)}', y_tag=None, x_log=False, y_log=False, x_range=None, y_range=None,
           colors=None,
           export=False, path=None, filename=None, lw=3, close=False):
    linewidth = lw
    
    if not close:
        app = qt.mkQApp()
    
    ph = qt.plot()
    ph.showGrid(x=True,y=True)
    if x_range is not None:
        ph.setXRange(*x_range)
    else:
        ph.setXRange(min([min(sp) for sp in spaces]), max([max(sp) for sp in spaces]))
    if y_range is not None:
        ph.setYRange(*y_range)
    else:
        ph.setYRange(min([min(v) for v in vals]), max([max(v) for v in vals]))
    ph.setLogMode(x=x_log, y=y_log)
    
    ph.setTitle(title=titletext, offset=(1000,500))#.format(dim))
    ph.setLabel('left', y_tag)
    ph.setLabel('bottom', x_tag)
    #ph.setXRange(0, 0.15)
    
    if names is not None:
        ph.addLegend(offset=(400, 30))
    
    
    #s = ph.plot(np.linspace(0.01,0.32,32), eps_max_freq0, title=sp_Title, pen='w')
    #s = ph.plot(np.linspace(0.01,0.32,32), eps_max_freq1, title=sp_Title, pen='w')
    
    
    if colors=='rgb':
        colors=[__get_color(fac/(len(vals)-1)) for fac in range(len(vals))]
    elif colors is None:
        colors=['r', 'g', 'b', 'y', 'c', 'm', 'w',
                (100,100,0), (0,100,100), (100,0,100),
                (100,200,0), (0,100,200), (200,0,100),
                (200,100,0), (0,200,100), (100,0,200)]
    else:
        colors = colors[::-1]
    
    
    for i in range(len(vals)-1,-1,-1):
        try:
            s = ph.plot(spaces[i],
                        vals[i], 
                        name=names[i] if names is not None else None,
                        pen=qt.mkPen(colors[i],
                        width=linewidth))
        except:
            print('Failed plotting '+names[i])
            raise
            
    #title_item = qt.TextItem(text=titletext, anchor=(0.5, 7), color='grey')
    #ph.addItem(title_item)
    #title_item.setPos(ph.getViewBox().viewRect().center())
    #font = QFont()
    #font.setPointSize(14)  # Adjust the font size as needed
    #title_item.setFont(font)

    if export:
        new_folder(path)
        
        exporter = qt.exporters.ImageExporter(ph.plotItem)
    
        # set export parameters if needed
        exporter.parameters()['width'] = 1200   # (note this also affects height parameter)
        
        # save to file
        exporter.export(path+filename)
        
        print(f'Saving to {path+filename}')
    
    def handleAppExit():
        # Add any cleanup tasks here
        print("closing")
    
    if close:
        ph.close()
    else:
        app.aboutToQuit.connect(handleAppExit)
        app.exec_()