Commit 245d4985 authored by Kelvin Loh's avatar Kelvin Loh

implement plotting backend selection

parent f8b0e474
......@@ -18,6 +18,7 @@
import warnings
from math import sqrt, pi
import numpy as np
from enum import Enum
try:
import matplotlib
......@@ -34,11 +35,34 @@ try:
warnings.warn("3D plotting not available.", RuntimeWarning)
has3d = False
except ImportError:
warnings.warn("matplotlib is not available, only iterator-providing "
"functions will work.", RuntimeWarning)
warnings.warn("matplotlib is not available, if other backends are "
"unavailable, only iterator-providing functions will work",
RuntimeWarning)
mpl_available = False
try:
import plotly.offline as plotly
import plotly.graph_objs as plotly_graph_objs
plotly.init_notebook_mode(connected=True)
plotly_available = True
except ImportError:
warnings.warn("plotly is not available, if other backends are unavailable,"
" only iterator-providing functions will work",
RuntimeWarning)
plotly_available = False
class Backends(Enum):
matplotlib = 0
plotly = 1
backend = Backends.matplotlib
if not ((mpl_available) or (plotly_available)):
backend = None
# Collections that allow for symbols and linewiths to be given in data space
# (not for general use, only implement what's needed for plotter)
def isarray(var):
......
......@@ -30,7 +30,8 @@ from . import system, builder, _common
from ._common import deprecate_args
__all__ = ['plot', 'map', 'bands', 'spectrum', 'current', 'density',
__all__ = ['set_backend',
'plot', 'map', 'bands', 'spectrum', 'current', 'density',
'interpolate_current', 'interpolate_density',
'streamplot', 'scalarplot',
'sys_leads_sites', 'sys_leads_hoppings', 'sys_leads_pos',
......@@ -41,6 +42,43 @@ __all__ = ['plot', 'map', 'bands', 'spectrum', 'current', 'density',
_p = _common.lazy_import('_plotter')
def set_backend(backend):
"""Set the plotting backend to use.
Parameters
----------
backend : str
Options are: 'matplotlib', 'plotly'.
"""
if ((_p.mpl_available) or (_p.plotly_available)):
try:
_p.backend = _p.Backends[backend]
except KeyError:
list_available_backends = []
if (_p.mpl_available):
list_available_backends.append('matplotlib')
if (_p.plotly_available):
list_available_backends.append('plotly')
error_message = "Tried to set an unknown backend \'{}\'.".format(
backend)
error_message += " Supported backends are {}".format(
list_available_backends)
raise RuntimeError(error_message)
else:
warnings.warn("Tried to set \'{}\' but is not "
"available.".format(backend), RuntimeWarning)
if (_p.is_ipython_kernel):
if ((_p.backend == _p.Backends.plotly) and
(not _p.init_notebook_mode_set)):
_p.init_notebook_mode_set = True
_p.plotly_module.init_notebook_mode(connected=True)
def get_backend():
return _p.backend
def _sample_array(array, n_samples, rng=None):
rng = _common.ensure_rng(rng)
la = len(array)
......@@ -106,13 +144,19 @@ def _maybe_output_fig(fig, file=None, show=True):
if fig is None:
return
if file is not None:
fig.canvas.print_figure(file, dpi=fig.dpi)
elif show:
# If there was no file provided, pyplot should already be available and
# we can import it safely without additional warnings.
from matplotlib import pyplot
pyplot.show()
if get_backend() == _p.Backends.matplotlib:
if file is not None:
fig.canvas.print_figure(file, dpi=fig.dpi)
elif show:
# If there was no file provided, pyplot should already be available
# and we can import it safely without additional warnings.
from matplotlib import pyplot
pyplot.show()
elif get_backend() == _p.Backends.plotly:
if file is not None:
_p.plotly.plot(fig, show_link=False, filename=file, auto_open=False)
if show:
_p.plotly.iplot(fig)
def set_colors(color, collection, cmap, norm=None):
......
......@@ -585,7 +585,8 @@ def main():
'tinyarray >= 1.2'],
extras_require={
# The oldest versions between: Debian stable, Ubuntu LTS
'plotting': 'matplotlib >= 2.1.1',
'plotting': ['matplotlib >= 2.1.1',
'plotly >= 2.2.2'],
'continuum': 'sympy >= 1.1.1',
# qsymm is only packaged on PyPI
'qsymm': 'qsymm >= 1.2.6',
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment