plot_network: Split plotting into several functions to make sphinx happy

Also addresses part of issue #50: plot_map works again with cartopy.

scripts/plot_network.py

@@ -15,27 +15,6 @@ Description
 
 """
 
-# Dirty work-around so that sphinx can import this module and get the
-# doc-string should be refactored in the style of the other scripts, ideally
-# several functions for the different plots
-if __name__ != "__main__":
-    import sys
-    sys.exit(0)
-
-if 'snakemake' not in globals():
-    from vresutils.snakemake import MockSnakemake, Dict
-    from snakemake.rules import expand
-    import yaml
-    snakemake = Dict()
-    snakemake = MockSnakemake(
-        path='..',
-        wildcards=dict(network='elec', simpl='', clusters='90', lv='1.25', opts='Co2L-3H', attr='p_nom', ext="pdf"),
-        input=dict(network="results/networks/{network}_s{simpl}_{clusters}_lv{lv}_{opts}.nc",
-                   tech_costs="data/costs.csv"),
-        output=dict(only_map="results/plots/{network}_s{simpl}_{clusters}_lv{lv}_{opts}_{attr}.{ext}",
-                    ext="results/plots/{network}_s{simpl}_{clusters}_lv{lv}_{opts}_{attr}_ext.{ext}")
-    )
-
 import pypsa
 
 from _helpers import load_network, aggregate_p, aggregate_costs
@@ -50,7 +29,9 @@ from itertools import product, chain
 from six.moves import map, zip
 from six import itervalues, iterkeys
 from collections import OrderedDict as odict
+import logging
 
+import cartopy.crs as ccrs
 import matplotlib.pyplot as plt
 import matplotlib as mpl
 from matplotlib.patches import Circle, Ellipse
@@ -83,6 +64,7 @@ def make_handler_map_to_scale_circles_as_in(ax, dont_resize_actively=False):
 def make_legend_circles_for(sizes, scale=1.0, **kw):
     return [Circle((0,0), radius=(s/scale)**0.5, **kw) for s in sizes]
 
+def set_plot_style():
     plt.style.use(['classic', 'seaborn-white',
                 {'axes.grid': False, 'grid.linestyle': '--', 'grid.color': u'0.6',
                     'hatch.color': 'white',
@@ -94,20 +76,16 @@ plt.style.use(['classic', 'seaborn-white',
                     # 'font.family': 'Times New Roman'
                 }])
 
-opts = snakemake.config['plotting']
+def plot_map(n, ax=None, attribute='p_nom', opts={}):
-map_figsize = opts['map']['figsize']
+    if ax is None:
-map_boundaries = opts['map']['boundaries']
+        ax = plt.gca()
-
-n = load_network(snakemake.input.network, snakemake.input.tech_costs, snakemake.config)
-
-scenario_opts = snakemake.wildcards.opts.split('-')
 
     ## DATA
     line_colors = {'cur': "purple",
                    'exp': to_rgba("red", 0.7)}
     tech_colors = opts['tech_colors']
 
-if snakemake.wildcards.attr == 'p_nom':
+    if attribute == 'p_nom':
         # bus_sizes = n.generators_t.p.sum().loc[n.generators.carrier == "load"].groupby(n.generators.bus).sum()
         bus_sizes = pd.concat((n.generators.query('carrier != "load"').groupby(['bus', 'carrier']).p_nom_opt.sum(),
                                n.storage_units.groupby(['bus', 'carrier']).p_nom_opt.sum()))
@@ -124,24 +102,23 @@ dict(Line=(line_widths_cur['Line'] / n.lines.s_nom > 1e-3)
         .map({True: line_colors['cur'], False: to_rgba(line_colors['cur'], 0.)}))
 
     ## FORMAT
-linewidth_factor = opts['map'][snakemake.wildcards.attr]['linewidth_factor']
+    linewidth_factor = opts['map'][attribute]['linewidth_factor']
-bus_size_factor  = opts['map'][snakemake.wildcards.attr]['bus_size_factor']
+    bus_size_factor  = opts['map'][attribute]['bus_size_factor']
 
     ## PLOT
-fig, ax = plt.subplots(figsize=map_figsize)
     n.plot(line_widths=pd.concat(line_widths_exp)/linewidth_factor,
            line_colors=dict(Line=line_colors['exp'], Link=line_colors['exp']),
            bus_sizes=bus_sizes/bus_size_factor,
            bus_colors=tech_colors,
            boundaries=map_boundaries,
-       basemap=True,
+           geomap=True,
            ax=ax)
     n.plot(line_widths=pd.concat(line_widths_cur)/linewidth_factor,
            line_colors=pd.concat(line_colors_with_alpha),
            bus_sizes=0,
            bus_colors=tech_colors,
            boundaries=map_boundaries,
-       basemap=False,
+           geomap=True,   # TODO : Turn to False, after the release of PyPSA 0.14.2 (refer to https://github.com/PyPSA/PyPSA/issues/75)
            ax=ax)
     ax.set_aspect('equal')
     ax.axis('off')
@@ -152,6 +129,7 @@ ax.axis('off')
 
 
     # Rasterize basemap
+    # TODO : Check if this also works with cartopy
     for c in ax.collections[:2]: c.set_rasterized(True)
 
     # LEGEND
@@ -200,15 +178,16 @@ for t in techs:
     l3 = ax.legend(handles, labels, loc="upper center",  bbox_to_anchor=(0.5, -0.), # bbox_to_anchor=(0.72, -0.05),
                 handletextpad=0., columnspacing=0.5, ncol=4, title='Technology')
 
-
+    return fig
-fig.savefig(snakemake.output.only_map, dpi=150,
-            bbox_inches='tight', bbox_extra_artists=[l1,l2,l3])
 
 #n = load_network(snakemake.input.network, opts, combine_hydro_ps=False)
 
-## Add total energy p
 
-ax1 = ax = fig.add_axes([-0.115, 0.625, 0.2, 0.2])
+def plot_total_energy_pie(n, ax=None):
+    """Add total energy pie plot"""
+    if ax is None:
+        ax = plt.gca()
+
     ax.set_title('Energy per technology', fontdict=dict(fontsize="medium"))
 
     e_primary = aggregate_p(n).drop('load', errors='ignore').loc[lambda s: s>0]
@@ -224,10 +203,11 @@ for t1, t2, i in zip(texts, autotexts, e_primary.index):
             t1.remove()
             t2.remove()
 
-## Add average system cost bar plot
+def plot_total_cost_bar(n, ax=None):
-# ax2 = ax = fig.add_axes([-0.1, 0.2, 0.1, 0.33])
+    """Add average system cost bar plot"""
-# ax2 = ax = fig.add_axes([-0.1, 0.15, 0.1, 0.37])
+    if ax is None:
-ax2 = ax = fig.add_axes([-0.075, 0.1, 0.1, 0.45])
+        ax = plt.gca()
+
     total_load = (n.snapshot_weightings * n.loads_t.p.sum(axis=1)).sum()
 
     def split_costs(n):
@@ -271,6 +251,46 @@ ax.set_xlim([0,1])
     ax.set_xticklabels([]) #["w/o\nEp", "w/\nEp"])
     ax.grid(True, axis="y", color='k', linestyle='dotted')
 
+
+if __name__ == "__main__":
+    if 'snakemake' not in globals():
+        from vresutils.snakemake import MockSnakemake, Dict
+        from snakemake.rules import expand
+
+        snakemake = Dict()
+        snakemake = MockSnakemake(
+            path='..',
+            wildcards=dict(network='elec', simpl='', clusters='90', lv='1.25', opts='Co2L-3H', attr='p_nom', ext="pdf"),
+            input=dict(network="results/networks/{network}_s{simpl}_{clusters}_lv{lv}_{opts}.nc",
+                    tech_costs="data/costs.csv"),
+            output=dict(only_map="results/plots/{network}_s{simpl}_{clusters}_lv{lv}_{opts}_{attr}.{ext}",
+                        ext="results/plots/{network}_s{simpl}_{clusters}_lv{lv}_{opts}_{attr}_ext.{ext}")
+        )
+
+    logging.basicConfig(level=snakemake.config['logging_level'])
+
+    set_plot_style()
+
+    opts = snakemake.config['plotting']
+    map_figsize = opts['map']['figsize']
+    map_boundaries = opts['map']['boundaries']
+
+    n = load_network(snakemake.input.network, snakemake.input.tech_costs, snakemake.config)
+
+    scenario_opts = snakemake.wildcards.opts.split('-')
+
+    fig, ax = plt.subplots(figsize=map_figsize, subplot_kw={"projection": ccrs.PlateCarree()})
+    plot_map(n, ax, snakemake.wildcards.attr, opts)
+
+    fig.savefig(snakemake.output.only_map, dpi=150,
+                bbox_inches='tight', bbox_extra_artists=[l1,l2,l3])
+
+    ax1 = fig.add_axes([-0.115, 0.625, 0.2, 0.2])
+    plot_total_energy_pie(n, ax1)
+
+    ax2 = fig.add_axes([-0.075, 0.1, 0.1, 0.45])
+    plot_total_cost_bar(n, ax2)
+
     #fig.tight_layout()
 
     ll = snakemake.wildcards.ll