# -*- coding: utf-8 -*-

#

# Licensed under the terms of the BSD 3-Clause

# (see plotpy/LICENSE for details)

# pylint: disable=C0103

"""

Interactive plot

----------------

The `pyplot` module provides an interactive plotting interface similar to

`Matplotlib`'s, i.e. with MATLAB-like syntax.

"""

import sys

import guidata

import numpy as np

from guidata.configtools import get_icon

from guidata.env import execenv

from guidata.qthelpers import is_qobject_valid, win32_fix_title_bar_background

from qtpy import QtCore as QC

from qtpy import QtGui as QG

from qtpy import QtWidgets as QW

from qtpy.QtPrintSupport import QPrinter

from plotpy import io

from plotpy._scaler import INTERP_AA, INTERP_LINEAR, INTERP_NEAREST

from plotpy.builder import make

from plotpy.config import _

from plotpy.mathutils.colormap import ALL_COLORMAPS

from plotpy.panels.contrastadjustment import ContrastAdjustment

from plotpy.panels.csection.cswidget import XCrossSection, YCrossSection

from plotpy.panels.itemlist import PlotItemList

from plotpy.plot.base import BasePlot, BasePlotOptions

from plotpy.plot.manager import PlotManager

_qapp = None

_interactive = False

_figures = {}

_current_fig = None

_current_axes = None

def create_qapplication(exec_loop=False):

"""Creating Qt application (only once) and eventually exec Qt main loop"""

global _qapp # pylint: disable=global-statement

if _qapp is None:

_qapp = guidata.qapplication()

if exec_loop and not execenv.unattended:

_qapp.exec()

class Window(QW.QMainWindow):

"""Figure window"""

def __init__(self, title):

super().__init__()

win32_fix_title_bar_background(self)

self.default_tool = None

self.plots = []

self.itemlist = PlotItemList(None)

self.contrast = ContrastAdjustment(None)

self.xcsw = XCrossSection(None)

self.ycsw = YCrossSection(None)

self.manager = PlotManager(self)

self.toolbar = QW.QToolBar(_("Tools"), self)

self.manager.add_toolbar(self.toolbar, "default")

self.toolbar.setMovable(True)

self.toolbar.setFloatable(True)

self.addToolBar(QC.Qt.TopToolBarArea, self.toolbar)

frame = QW.QFrame(self)

self.setCentralWidget(frame)

self.layout = QW.QGridLayout()

layout = QW.QVBoxLayout(frame)

frame.setLayout(layout)

layout.addLayout(self.layout)

self.frame = frame

self.setWindowTitle(title)

self.setWindowIcon(get_icon("plotpy.svg"))

def closeEvent(self, event):

"""Reimplementing QWidget.closeEvent

:param event:

"""

global _figures, _current_fig, _current_axes

figure_title = str(self.windowTitle())

if _figures.pop(figure_title) == _current_fig:

_current_fig = None

_current_axes = None

for panel in (self.itemlist, self.contrast, self.xcsw, self.ycsw):

if is_qobject_valid(panel):

panel.close()

event.accept()

def add_plot(self, i, j, plot):

"""Add plot to window

:param i:

:param j:

:param plot:

"""

i_int = int(i)

j_int = int(j)

self.layout.addWidget(plot, i_int, j_int)

self.manager.add_plot(plot)

self.plots.append(plot)

def replot(self):

"""Replot all plots"""

for plot in self.plots:

plot.replot()

item = plot.get_default_item()

if item is not None:

plot.set_active_item(item)

item.unselect()

def add_panels(self, images=False):

"""Add panels

:param images:

"""

self.manager.add_panel(self.itemlist)

if images:

for panel in (self.ycsw, self.xcsw, self.contrast):

panel.hide()

self.manager.add_panel(panel)

def register_tools(self, images=False):

"""Register tools

:param images:

"""

if images:

self.manager.register_all_image_tools()

else:

self.manager.register_all_curve_tools()

def display(self):

"""Display window"""

self.show()

self.replot()

self.manager.get_default_tool().activate()

self.manager.update_tools_status()

class Figure:

"""Object representing a figure"""

def __init__(self, title):

self.axes = {}

self.title = title

self.win = None

def get_axes(self, i, j):

"""Get axes

:param i:

:param j:

:return:

"""

if (i, j) in self.axes:

return self.axes[(i, j)]

ax = Axes()

self.axes[(i, j)] = ax

return ax

def build_window(self):

"""Build window"""

create_qapplication()

self.win = Window(title=self.title)

images = False

for (i, j), ax in list(self.axes.items()):

ax.setup_window(i, j, self.win)

if ax.images:

images = True

self.win.add_panels(images=images)

self.win.register_tools(images=images)

def show(self):

"""Show window"""

if not self.win:

self.build_window()

self.win.display()

def save(self, fname, format, draft):

"""Save figure

:param fname:

:param format:

:param draft:

"""

if isinstance(fname, str):

if format == "pdf":

try:

mode = QPrinter.HighResolution

except AttributeError:

# Some PySide6 / PyQt6 versions do not have this attribute on Linux

mode = QPrinter.ScreenResolution

if draft:

mode = QPrinter.ScreenResolution

printer = QPrinter(mode)

try:

printer.setOutputFormat(QPrinter.PdfFormat)

except AttributeError:

# PyQt6 on Linux

printer.setPrinterName("")

printer.setPageOrientation(QG.QPageLayout.Landscape)

printer.setOutputFileName(fname)

printer.setCreator("plotpy.pyplot")

self.print_(printer)

else:

if self.win is None:

self.show()

pixmap = self.win.centralWidget().grab()

pixmap.save(fname, format.upper())

else:

# Buffer

fd = fname

assert hasattr(fd, "write"), "object is not file-like as expected"

if self.win is None:

self.show()

pixmap = self.win.centralWidget().grab()

buff = QC.QBuffer()

buff.open(QC.QIODevice.ReadWrite)

pixmap.save(buff, format.upper())

fd.write(buff.data())

buff.close()

fd.seek(0)

def print_(self, device):

"""Print figure

:param device:

"""

if not self.win:

self.build_window()

W = device.width()

H = device.height()

coords = np.array(list(self.axes.keys()))

imin = coords[:, 0].min()

imax = coords[:, 0].max()

jmin = coords[:, 1].min()

jmax = coords[:, 1].max()

w = W / (jmax - jmin + 1)

h = H / (imax - imin + 1)

paint = QG.QPainter(device)

for (i, j), ax in list(self.axes.items()):

oy = (i - imin) * h

ox = (j - jmin) * w

ax.widget.print_(paint, QC.QRect(ox, oy, w, h))

class Axes:

"""Object handling plots, images and axes properties"""

def __init__(self):

self.plots = []

self.images = []

self.last = None

self.legend_position = None

self.grid = False

self.xlabel = ("", "")

self.ylabel = ("", "")

self.xcolor = ("black", "black") # axis label colors

self.ycolor = ("black", "black") # axis label colors

self.zlabel = None

self.yreverse = False

self.colormap = "jet"

self.xscale = "lin"

self.yscale = "lin"

self.xlimits = None

self.ylimits = None

self.widget = None

self.main_widget = None

def add_legend(self, position):

"""Add legend

:param position:

"""

self.legend_position = position

def set_grid(self, grid):

"""Set grid

:param grid:

"""

self.grid = grid

def set_xlim(self, xmin, xmax):

"""Set x limits

:param xmin:

:param xmax:

"""

self.xlimits = xmin, xmax

self._update_plotwidget()

def set_ylim(self, ymin, ymax):

"""Set y limits

:param ymin:

:param ymax:

"""

self.ylimits = ymin, ymax

self._update_plotwidget()

def add_plot(self, item):

"""Add plot

:param item:

"""

self.plots.append(item)

self.last = item

def add_image(self, item):

"""Add image

:param item:

"""

self.images.append(item)

self.last = item

def setup_window(self, i, j, win):

"""Setup window

:param i:

:param j:

:param win:

"""

if self.images:

plot = self.setup_image(i, j, win)

else:

plot = self.setup_plot(i, j, win)

self.widget = plot

plot.do_autoscale()

self._update_plotwidget()

def _update_plotwidget(self):

"""Update plot widget"""

p = self.main_widget

if p is None:

return

if self.grid:

p.gridparam.maj_xenabled = True

p.gridparam.maj_yenabled = True

p.gridparam.update_grid(p)

p.set_axis_color("bottom", self.xcolor[0])

p.set_axis_color("top", self.xcolor[1])

p.set_axis_color("left", self.ycolor[0])

p.set_axis_color("right", self.ycolor[1])

if self.xlimits is not None:

p.set_axis_limits("bottom", *self.xlimits)

if self.ylimits is not None:

p.set_axis_limits("left", *self.ylimits)

def setup_image(self, i, j, win: Window):

"""Setup image

:param i:

:param j:

:param win:

:return:

"""

options = BasePlotOptions(

xlabel=self.xlabel,

ylabel=self.ylabel,

yreverse=self.yreverse,

type="image",

)

p = BasePlot(win, options=options)

self.main_widget = p

win.add_plot(i, j, p)

for item in self.images + self.plots:

if item in self.images:

item.set_color_map(self.colormap)

p.add_item(item)

if self.legend_position is not None:

p.add_item(make.legend(self.legend_position))

return p

def setup_plot(self, i, j, win):

"""Setup plot

:param i:

:param j:

:param win:

:return:

"""

options = BasePlotOptions(xlabel=self.xlabel, ylabel=self.ylabel, type="curve")

p = BasePlot(win, options=options)

self.main_widget = p

win.add_plot(i, j, p)

for item in self.plots:

p.add_item(item)

p.enable_used_axes()

active_item = p.get_active_item(force=True)

p.set_scales(self.xscale, self.yscale)

active_item.unselect()

if self.legend_position is not None:

p.add_item(make.legend(self.legend_position))

return p

def _make_figure_title(N=None):

"""Make figure title"""

global _figures

if N is None:

N = len(_figures) + 1

if isinstance(N, str):

return N

else:

return f"Figure {N:d}"

def figure(N=None):

"""Create a new figure"""

global _figures, _current_fig, _current_axes

title = _make_figure_title(N)

if title in _figures:

f = _figures[title]

else:

f = Figure(title)

_figures[title] = f

_current_fig = f

_current_axes = None

return f

def gcf():

"""Get current figure"""

global _current_fig

if _current_fig:

return _current_fig

else:

return figure()

def gca():

"""Get current axes"""

create_qapplication() # Necessary because of setIcon call during item creation

global _current_axes

if not _current_axes:

axes = gcf().get_axes(1, 1)

_current_axes = axes

return _current_axes

def show(mainloop=True):

"""

Show all figures and enter Qt event loop

This should be the last line of your script

"""

global _figures, _interactive, _current_fig

for fig in list(_figures.values()):

fig.show()

if not _interactive:

if not _current_fig:

print("Warning: must create a figure before showing it", file=sys.stderr)

elif mainloop:

create_qapplication(exec_loop=True)

def _show_if_interactive():

"""Show all figures if interactive mode is on"""

global _interactive

if _interactive:

show()

def subplot(n, m, k):

"""

Create a subplot command

Example::

import numpy as np

x = np.linspace(-5, 5, 1000)

figure(1)

subplot(2, 1, 1)

plot(x, np.sin(x), "r+")

subplot(2, 1, 2)

plot(x, np.cos(x), "g-")

show()

"""

global _current_axes

lig = (k - 1) // m

col = (k - 1) % m

fig = gcf()

axe = fig.get_axes(lig, col)

_current_axes = axe

return axe

def plot(*args, **kwargs):

"""

Plot curves

Example::

import numpy as np

x = np.linspace(-5, 5, 1000)

plot(x, np.sin(x), "r+")

plot(x, np.cos(x), "g-")

show()

"""

axe = gca()

curves = make.mcurve(*args, **kwargs)

if not isinstance(curves, list):

curves = [curves]

for curve in curves:

axe.add_plot(curve)

_show_if_interactive()

return curves

def plotyy(x1, y1, x2, y2):

"""

Plot curves with two different y axes

Example::

import numpy as np

x = np.linspace(-5, 5, 1000)

plotyy(x, np.sin(x), x, np.cos(x))

ylabel("sine", "cosine")

show()

"""

axe = gca()

curve1 = make.mcurve(x1, y1, yaxis="left")

curve2 = make.mcurve(x2, y2, yaxis="right")

axe.ycolor = (curve1.param.line.color, curve2.param.line.color)

axe.add_plot(curve1)

axe.add_plot(curve2)

_show_if_interactive()

return [curve1, curve2]

def hist(data, bins=None, logscale=None, title=None, color=None):

"""

Plot 1-D histogram

Example::

from numpy.random import normal

data = normal(0, 1, (2000, ))

hist(data)

show()

"""

axe = gca()

curve = make.histogram(

data, bins=bins, logscale=logscale, title=title, color=color, yaxis="left"

)

axe.add_plot(curve)

_show_if_interactive()

return [curve]

def semilogx(*args, **kwargs):

"""

Plot curves with logarithmic x-axis scale

Example::

import numpy as np

x = np.linspace(-5, 5, 1000)

semilogx(x, np.sin(12*x), "g-")

show()

"""

axe = gca()

axe.xscale = "log"

curve = make.mcurve(*args, **kwargs)

axe.add_plot(curve)

_show_if_interactive()

return [curve]

def semilogy(*args, **kwargs):

"""

Plot curves with logarithmic y-axis scale

Example::

import numpy as np

x = np.linspace(-5, 5, 1000)

semilogy(x, np.sin(12*x), "g-")

show()

"""

axe = gca()

axe.yscale = "log"

curve = make.mcurve(*args, **kwargs)

axe.add_plot(curve)

_show_if_interactive()

return [curve]

def loglog(*args, **kwargs):

"""

Plot curves with logarithmic x-axis and y-axis scales

Example::

import numpy as np

x = np.linspace(-5, 5, 1000)

loglog(x, np.sin(12*x), "g-")

show()

"""

axe = gca()

axe.xscale = "log"

axe.yscale = "log"

curve = make.mcurve(*args, **kwargs)

axe.add_plot(curve)

_show_if_interactive()

return [curve]

def errorbar(*args, **kwargs):

"""

Plot curves with error bars

Example::

import numpy as np

x = np.linspace(-5, 5, 1000)

errorbar(x, -1+x**2/20+.2*np.random.rand(len(x)), x/20)

show()

"""

axe = gca()

curve = make.merror(*args, **kwargs)

axe.add_plot(curve)

_show_if_interactive()

return [curve]

def imread(fname, to_grayscale=False):

"""Read data from *fname*"""

return io.imread(fname, to_grayscale=to_grayscale)

def imshow(data, interpolation=None, mask=None):

"""

Display the image in *data* to current axes

interpolation: 'nearest', 'linear' (default), 'antialiasing'

Example::

import numpy as np

x = np.linspace(-5, 5, 1000)

img = np.fromfunction(lambda x, y: np.sin((x/200.)*(y/200.)**2), (1000, 1000))

gray()

imshow(img)

show()

"""

axe = gca()

if isinstance(data, np.ma.MaskedArray) and mask is None:

mask = data.mask

data = data.data

if mask is None:

img = make.image(data)

else:

img = make.maskedimage(data, mask, show_mask=True)

if interpolation is not None:

interp_dict = {

"nearest": INTERP_NEAREST,

"linear": INTERP_LINEAR,

"antialiasing": INTERP_AA,

}

assert interpolation in interp_dict, "invalid interpolation option"

img.set_interpolation(interp_dict[interpolation], size=5)

axe.add_image(img)

axe.yreverse = True

_show_if_interactive()

return [img]

def pcolor(*args):

"""

Create a pseudocolor plot of a 2-D array

Example::

import numpy as np

r = np.linspace(1., 16, 100)

th = np.linspace(0., np.pi, 100)

R, TH = np.meshgrid(r, th)

X = R*np.cos(TH)

Y = R*np.sin(TH)

Z = 4*TH+R

pcolor(X, Y, Z)

show()

"""

axe = gca()

img = make.pcolor(*args)

axe.add_image(img)

axe.yreverse = len(args) == 1

_show_if_interactive()

return [img]

def interactive(state):

"""Toggle interactive mode"""

global _interactive

_interactive = state

def ion():

"""Turn interactive mode on"""

interactive(True)

def ioff():

"""Turn interactive mode off"""

interactive(False)

# TODO: The following functions (title, xlabel, ...) should update an already

# shown figure to be compatible with interactive mode -- for now it just

# works if these functions are called before showing the figure

def title(text):

"""Set current figure title"""

global _figures

fig = gcf()

_figures.pop(fig.title)

fig.title = text

_figures[text] = fig

def xlabel(bottom="", top=""):

"""Set current x-axis label"""

assert isinstance(bottom, str) and isinstance(top, str)

axe = gca()

axe.xlabel = (bottom, top)

def ylabel(left="", right=""):

"""Set current y-axis label"""

assert isinstance(left, str) and isinstance(right, str)

axe = gca()

axe.ylabel = (left, right)

def zlabel(label):

"""Set current z-axis label"""

assert isinstance(label, str)

axe = gca()

axe.zlabel = label

def yreverse(reverse):

"""

Set y-axis direction of increasing values

reverse = False (default)

y-axis values increase from bottom to top

reverse = True

y-axis values increase from top to bottom

"""

assert isinstance(reverse, bool)

axe = gca()

axe.yreverse = reverse

def grid(act):

"""Toggle grid visibility"""

axe = gca()

axe.set_grid(act)

def legend(pos="TR"):

"""Add legend to current axes (pos='TR', 'TL', 'BR', ...)"""

axe = gca()

axe.add_legend(pos)

def colormap(name):

"""Set color map to *name*"""

axe = gca()

axe.colormap = name

def _add_colormaps(glbs):

"""Add colormap functions to *glbs*"""

for cmap_name in ALL_COLORMAPS.keys():

glbs[cmap_name] = lambda name=cmap_name: colormap(name)

glbs[cmap_name].__doc__ = "Set color map to '%s'" % cmap_name

_add_colormaps(globals())

def close(N=None, all=False):

"""Close figure"""

global _figures, _current_fig, _current_axes

if all:

_figures = {}

_current_fig = None

_current_axes = None

return

if N is None:

fig = gcf()

else:

fig = figure(N)

fig.close()

def savefig(fname, format=None):

"""

Save figure

Currently supports QImageWriter formats only

(see https://doc.qt.io/qt-5/qimagewriter.html#supportedImageFormats)

"""

if not isinstance(fname, str) and format is None:

# Buffer/fd

format = "png"

if format is None:

format = fname.rsplit(".", 1)[-1].lower()

fmts = [fmt.data().decode() for fmt in QG.QImageWriter.supportedImageFormats()]

assert format in fmts, _(

"Function 'savefig' currently supports the following formats:\n%s"

) % ", ".join(fmts)

else:

format = format.lower()

fig = gcf()

fig.save(fname, format)