Brillplot: work with Agg backend, create 3d prettyplot base

It's important that an Axis passed to the Pymatgen Brillouin plotter
is already 3D or we get some rather confusing error messages.

small cleanup

Author: ajjackson

sumo/cli/brillplot.py

@@ -23,10 +23,11 @@
 
 from pymatgen.io.vasp.outputs import BSVasprun
 from pymatgen.electronic_structure.bandstructure import get_reconstructed_band_structure
-from pymatgen.electronic_structure.plotter import BSPlotter
+from pymatgen.electronic_structure.plotter import plot_brillouin_zone
 import matplotlib as mpl
-
 mpl.use("Agg")
+from sumo.plotting import pretty_plot_3d
+
 __author__ = "Arthur Youd"
 __version__ = "1.0"
 __maintainer__ = "Alex Ganose"
@@ -35,6 +36,7 @@
 
 
 def brillplot(filenames=None, prefix=None, directory=None,
+              width=6, height=6, fonts=None,
               image_format="pdf", dpi=400):
     """Generate plot of first brillouin zone from a band-structure calculation.
     Args:
@@ -57,17 +59,28 @@ def brillplot(filenames=None, prefix=None, directory=None,
         bs = vr.get_band_structure(line_mode=True)
         bandstructures.append(bs)
     bs = get_reconstructed_band_structure(bandstructures)
-    plotter = BSPlotter(bs)
-    plt = plotter.plot_brillouin()
+
+    labels = {}
+    for k in bs.kpoints:
+        if k.label:
+            labels[k.label] = k.frac_coords
+
+    lines = []
+    for b in bs.branches:
+        lines.append([bs.kpoints[b['start_index']].frac_coords,
+                      bs.kpoints[b['end_index']].frac_coords])
+
+    plt = pretty_plot_3d(width, height, dpi=dpi, fonts=fonts)
+    fig = plot_brillouin_zone(bs.lattice_rec, lines=lines, labels=labels,
+                              ax=plt.gca())
 
     basename = "brillouin.{}".format(image_format)
     filename = "{}_{}".format(prefix, basename) if prefix else basename
     if directory:
         filename = os.path.join(directory, filename)
-    plt.savefig(filename, format=image_format, dpi=dpi, bbox_inches="tight")
+    fig.savefig(filename, format=image_format, dpi=dpi, bbox_inches="tight")
     return plt
 
-
 def find_vasprun_files():
     """Search for vasprun files from the current directory.
 
@@ -99,7 +112,7 @@ def find_vasprun_files():
 
 def _get_parser():
     parser = argparse.ArgumentParser(description="""
-    brillplot is a script to produce publication-ready 
+    brillplot is a script to produce publication-ready
     brillouin zone diagrams""",
                                      epilog="""
     Author: {}

sumo/plotting/__init__.py

@@ -12,6 +12,7 @@
 from matplotlib.collections import LineCollection
 from matplotlib import rc, rcParams
 from pkg_resources import resource_filename
+from mpl_toolkits.mplot3d import Axes3D
 
 colour_cache = {}
 
@@ -38,7 +39,6 @@ def styled_plot(*style_sheets):
     """
 
     def decorator(get_plot):
-
         def wrapper(*args, fonts=None, style=None, no_base_style=False,
                     **kwargs):
 
@@ -96,6 +96,27 @@ def pretty_plot(width=None, height=None, plt=None, dpi=None):
     return plt
 
 
+def pretty_plot_3d(width=5, height=5, plt=None, dpi=None, fonts=None):
+    if plt is None:
+        plt = matplotlib.pyplot
+        if width is None:
+            width = matplotlib.rcParams['figure.figsize'][0]
+        if height is None:
+            height = matplotlib.rcParams['figure.figsize'][1]
+
+        if dpi is not None:
+            matplotlib.rcParams['figure.dpi'] = dpi
+
+        fig = plt.figure(figsize=(width, height), dpi=dpi)
+
+    else:
+        fig = plt.gcf()
+
+    ax = fig.add_subplot(111, projection='3d')
+
+    return plt
+
+
 def pretty_subplot(nrows, ncols, width=None, height=None, sharex=True,
                    sharey=True, dpi=None, plt=None, gridspec_kw=None):
     """Get a :obj:`matplotlib.pyplot` subplot object with pretty defaults.
@@ -133,7 +154,6 @@ def pretty_subplot(nrows, ncols, width=None, height=None, sharex=True,
         plt.subplots(nrows, ncols, sharex=sharex, sharey=sharey, dpi=dpi,
                      figsize=(width, height), facecolor='w',
                      gridspec_kw=gridspec_kw)
-
     return plt