Merge pull request #31 from AngelFP/add_2d_vis

Add 2D matplotlib visualizer

visualpic/helper_functions.py

@@ -81,13 +81,28 @@ def get_common_timesteps(data_list):
     --------
     An array containing only the common time steps.
     """
-    timesteps = np.array([])
-    for i, data_element in enumerate(data_list):
+    common_ts = np.array([])
+    ts_list = []
+    for data_element in data_list:
+        ts_list.append(data_element.timesteps)
+    if len(ts_list) > 0:
+        common_ts = get_common_array_elements(ts_list)
+    return common_ts
+
+
+def get_common_array_elements(array_list):
+    """
+    Returns an array containing only the common elements between all the
+    arrays in 'array_list'
+    
+    """
+    for i, array in enumerate(array_list):
         if i == 0:
-            timesteps = data_element.timesteps
+            common_els = array
         else:
-            timesteps = np.intersect1d(timesteps, data_element.timesteps)
-    return timesteps
+            common_els = np.intersect1d(common_els, array)
+    return common_els
+
 
 
 def get_closest_timestep(time_step, time_steps):

visualpic/ui/basic_plot_window.py

@@ -0,0 +1,134 @@
+"""
+This file is part of VisualPIC.
+
+The module contains the class for the basic Qt window for the matplotlib
+visualizer.
+
+Copyright 2016-2020, Angel Ferran Pousa.
+License: GNU GPL-3.0.
+"""
+
+
+from pkg_resources import resource_filename
+import platform
+import ctypes
+
+import numpy as np
+from PyQt5.Qt import Qt, QStyleFactory
+from PyQt5 import QtCore, QtWidgets, QtGui
+from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
+
+from visualpic.helper_functions import (
+    get_closest_timestep, get_next_timestep, get_previous_timestep)
+
+# code for proper scaling in high-DPI screens. Move this somewhere else once \
+# final UI is implemented.
+if platform.system() == 'Windows':
+    myappid = 'visualpic'  # arbitrary string
+    ctypes.windll.shell32.SetCurrentProcessExplicitAppUserModelID(myappid)
+    ctypes.windll.user32.SetProcessDPIAware()
+# Enable scaling for high DPI displays
+QtWidgets.QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)
+QtWidgets.QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
+QtWidgets.QApplication.setStyle(QStyleFactory.create('Fusion'))
+
+
+class BasicPlotWindow(QtWidgets.QMainWindow):
+
+    """Basic Qt window for interactive visualization matplotlib plots."""
+
+    def __init__(self, vp_figure, mpl_visualizer, parent=None):
+        super().__init__(parent=parent)
+        self.vp_figure = vp_figure
+        self.mpl_vis = mpl_visualizer
+        self.available_timesteps = self.vp_figure.get_available_timesteps()
+        self.timestep_change_callbacks = []
+        self.setup_interface()
+        self.register_ui_events()
+        self.show()
+
+    def setup_interface(self):
+        self.resize(600, 400)
+        self.setWindowTitle('VisualPIC - 2D Viewer')
+        icon_path = resource_filename('visualpic.ui.icons',
+                                      'vp_logo_sq_blue.png')
+        self.setWindowIcon(QtGui.QIcon(icon_path))
+        self.frame = QtWidgets.QFrame()
+        self.vl = QtWidgets.QVBoxLayout()
+        self.dialog_dict = {}
+        self.settings_dialog = None
+
+        self.figure_canvas = FigureCanvasQTAgg(self.vp_figure._mpl_figure)
+        self.vl.addWidget(self.figure_canvas)
+
+        self.hl = QtWidgets.QHBoxLayout()
+        self.timestep_line_edit = QtWidgets.QLineEdit()
+        self.timestep_line_edit.setReadOnly(True)
+        self.timestep_line_edit.setAlignment(Qt.AlignCenter)
+        self.timestep_line_edit.setText(str(self.vp_figure._current_timestep))
+        self.timestep_line_edit.setMaximumWidth(50)
+        self.hl.addWidget(self.timestep_line_edit)
+        self.prev_button = QtWidgets.QPushButton()
+        prev_icon_path = resource_filename('visualpic.ui.icons',
+                                           'left_arrow.png')
+        self.prev_button.setIcon(QtGui.QIcon(prev_icon_path))
+        if len(self.available_timesteps) == 0:
+            self.prev_button.setEnabled(False)
+        self.hl.addWidget(self.prev_button)
+        self.next_button = QtWidgets.QPushButton()
+        next_icon_path = resource_filename('visualpic.ui.icons',
+                                           'right_arrow.png')
+        self.next_button.setIcon(QtGui.QIcon(next_icon_path))
+        if len(self.available_timesteps) == 0:
+            self.next_button.setEnabled(False)
+        self.hl.addWidget(self.next_button)
+        self.timestep_slider = QtWidgets.QSlider(Qt.Horizontal)
+        if len(self.available_timesteps) > 0:
+            self.timestep_slider.setRange(np.min(self.available_timesteps),
+                                          np.max(self.available_timesteps))
+            self.timestep_slider.setValue(self.vp_figure._current_timestep)
+        else:
+            self.timestep_slider.setEnabled(False)
+        self.hl.addWidget(self.timestep_slider)
+
+
+        self.vl.addLayout(self.hl)
+        self.frame.setLayout(self.vl)
+        self.setCentralWidget(self.frame)
+
+    def register_ui_events(self):
+        self.timestep_slider.sliderReleased.connect(
+            self.timestep_slider_released)
+        self.timestep_slider.valueChanged.connect(
+            self.timestep_slider_value_changed)
+        self.prev_button.clicked.connect(self.prev_button_clicked)
+        self.next_button.clicked.connect(self.next_button_clicked)
+
+    def prev_button_clicked(self):
+        current_ts = self.timestep_slider.value()
+        prev_ts = get_previous_timestep(current_ts, self.available_timesteps)
+        if prev_ts != current_ts:
+            self.timestep_slider.setValue(prev_ts)
+        self.vp_figure.generate(prev_ts)
+        self.figure_canvas.draw()
+        # self.render_timestep(prev_ts)
+
+    def next_button_clicked(self):
+        current_ts = self.timestep_slider.value()
+        next_ts = get_next_timestep(current_ts, self.available_timesteps)
+        if next_ts != current_ts:
+            self.timestep_slider.setValue(next_ts)
+        self.vp_figure.generate(next_ts)
+        self.figure_canvas.draw()
+        # self.render_timestep(next_ts)
+
+    def timestep_slider_released(self):
+        value = self.timestep_slider.value()
+        value = get_closest_timestep(value, self.available_timesteps)
+        self.timestep_slider.setValue(value)
+        self.vp_figure.generate(value)
+        self.figure_canvas.draw()
+        # self.render_timestep(value)
+
+    def timestep_slider_value_changed(self, value):
+        self.timestep_line_edit.setText(str(value))

visualpic/visualization/matplotlib/mpl_visualizer.py

@@ -0,0 +1,150 @@
+"""
+This file is part of VisualPIC.
+
+The module contains the classes for 2D visualization with matplotlib.
+
+Copyright 2016-2020, Angel Ferran Pousa.
+License: GNU GPL-3.0.
+"""
+
+import sys
+
+from PyQt5 import QtWidgets
+
+from visualpic.ui.basic_plot_window import BasicPlotWindow
+from .plot_containers import VPFigure, FieldSubplot, ParticleSubplot
+from .rc_params import rc_params, rc_params_dark
+
+
+class MplVisualizer():
+    """Class taking care of visualizing the data with matplotlib."""
+
+    def __init__(self):
+        self._figure_list = []
+        self._current_figure = None
+
+    def figure(self, fig_idx=None):
+        """Set current figure in which to plot.
+
+        Parameters
+        ----------
+        fig_idx : int, optional
+            Index of the figure, by default None
+
+        Returns
+        -------
+        VPFigure
+        """
+        if fig_idx is not None:
+            n_figs = len(self._figure_list)
+            if n_figs < fig_idx + 1:
+                for i in range(fig_idx + 1 - n_figs):
+                    self._figure_list.append(VPFigure(rc_params=rc_params))
+            fig = self._figure_list[fig_idx]
+        else:
+            fig = VPFigure(rc_params=rc_params)
+            self._figure_list.append(fig)
+        self._set_current_figure(fig)
+        return fig
+
+    def field_plot(
+            self, field, field_units=None, axes_units=None, time_units=None,
+            slice_dir=None, slice_pos=0.5, m='all', theta=0., vmin=None,
+            vmax=None, cmap=None, stacked=True, cbar=True):
+        """Add a field plot to the figure.
+
+        Parameters
+        ----------
+        field : Field or list of Fields
+            The VisualPIC field(s) to plot.
+        field_units : str, optional
+            Desired field units, by default None
+        axes_units : str, optional
+            Desired axes units, by default None
+        time_units : str, optional
+            Desired time units, by default None
+        slice_dir : str, optional
+            Direction along which to slice the field, by default None
+        slice_pos : float, optional
+            Location of the slice in the `slice_dir` axes. Value between
+            0 an 1, by default 0.5.
+        m : str or int, optional
+            For fields in `thetaMode` geometry, azimuthal mode to show, by
+            default 'all'.
+        theta : float, optional
+            For fields in `thetaMode` geometry, angle at which to show the
+            fields, by default 0.
+        vmin : float or list of floats, optional
+            Minimum of the colormap range.
+        vmax : float or list of floats, optional
+            Maximum of the colormap range.
+        cmap : str or matplotlib Colormap, or a list of them.
+            The colormap to use for each field.
+        stacked : bool, optional
+            Whether to stack the all fields in the same axes (on top of each
+            other) or show them in individual axes, by default True.
+        cbar : bool, optional
+            Whether to include a colorbar, by default True.
+        """
+        fig = self._get_current_figure()
+        subplot = FieldSubplot(
+            field, field_units=field_units, axes_units=axes_units,
+            time_units=time_units, slice_dir=slice_dir, slice_pos=slice_pos,
+            m=m, theta=theta, vmin=vmin, vmax=vmax, cmap=cmap, stacked=stacked,
+            cbar=cbar)
+        fig.add_subplot(subplot)
+
+    def particle_plot(
+            self, species, x='x', y='y', x_units=None, y_units=None,
+            q_units=None, time_units=None, cbar=True):
+        """Add a particle plot to the figure.
+
+        Parameters
+        ----------
+        species : Species
+            The VisualPIC species to show.
+        x : str, optional
+            Name of the coordinate plotted in the x axis, by default 'x'
+        y : str, optional
+            Name of the coordinate plotted in the y axis, by default 'y'
+        x_units : str, optional
+            Units of the x coordinate
+        y_units : str, optional
+            Units of the x coordinate
+        q_units : str, optional
+            Units of the particle charge
+        time_units : str, optional
+            Units of the time.
+        cbar : bool, optional
+            Whether to show a colorbar, by default True
+        """
+        fig = self._get_current_figure()
+        subplot = ParticleSubplot(
+            species, x=x, y=y, x_units=x_units, y_units=y_units,
+            q_units=q_units, time_units=time_units, cbar=cbar)
+        fig.add_subplot(subplot)
+
+    def show(self, timestep=0):
+        """Show figure(s).
+
+        Parameters
+        ----------
+        timestep : int, optional
+            Time step at which to show the data, by default 0
+        """
+        app = QtWidgets.QApplication(sys.argv)
+        for figure in self._figure_list:
+            figure.generate(timestep)
+        self.windows = []
+        for figure in self._figure_list:
+            self.windows.append(BasicPlotWindow(figure, self))
+        app.exec_()
+
+    def _set_current_figure(self, figure):
+        self._current_figure = figure
+
+    def _get_current_figure(self):
+        if self._current_figure is None:
+            return self.figure()
+        else:
+            return self._current_figure

visualpic/visualization/matplotlib/plot_containers/__init__.py

@@ -0,0 +1,5 @@
+from .figure import VPFigure
+from .subplots import FieldSubplot, ParticleSubplot
+
+
+__all__ = ['VPFigure', 'FieldSubplot', 'ParticleSubplot']

visualpic/visualization/matplotlib/plot_containers/figure.py

@@ -0,0 +1,34 @@
+import matplotlib.pyplot as plt
+from matplotlib.figure import Figure
+
+from visualpic.helper_functions import get_common_array_elements
+
+
+class VPFigure():
+    """
+    Class defining a VisualPIC figure. Mainly a wrapper around a matplotlib
+    Figure.
+    """
+    def __init__(self, rc_params={}, **kwargs):
+        self._subplot_list = []
+        with plt.rc_context(rc_params):
+            self._mpl_figure = Figure(tight_layout=True)
+        self._current_timestep = -1
+
+    def add_subplot(self, subplot):
+        self._subplot_list.append(subplot)
+
+    def generate(self, timestep):
+        self._mpl_figure.clear()
+        gs = self._mpl_figure.add_gridspec(1, len(self._subplot_list))
+        self._current_timestep = timestep
+        for i, subplot in enumerate(self._subplot_list):
+            # ax = self._mpl_figure.add_subplot(gs[i])
+            subplot.plot(timestep, gs[i], self._mpl_figure)
+        # self._mpl_figure.tight_layout()
+
+    def get_available_timesteps(self):
+        ts_list = []
+        for subplot in self._subplot_list:
+            ts_list.append(subplot.get_available_timesteps())
+        return get_common_array_elements(ts_list)

visualpic/visualization/matplotlib/plot_containers/subplots/__init__.py

@@ -0,0 +1,5 @@
+from .field_subplot import FieldSubplot
+from .particle_subplot import ParticleSubplot
+
+
+__all__ = ['FieldSubplot', 'ParticleSubplot']