Several field subplots on same figure

The plot option of the field subcommand now uses the same logic
as rprof and time to describe the list of variables to plot.

The shrinkcb option of the field subcommand is removed since the
size of colorbars is now automatically determined.

Author: amorison

docs/sources/cookbook/field.rst

@@ -17,20 +17,44 @@ will plot the last snapshot of the temperature field.
 
 ::
 
-   % stagpy field -s : -o T,v3
+    % stagpy field -s : -o T-v3
 
 will plot all the snapshots of the temperature and the vertical/radial velocity
 on separate figures.
 
 ::
 
-   % stagpy field -s 3:8:2 -o T+stream
+    % stagpy field -s 3:8:2 -o T,stream
 
 will plot the temperature field with isocontours of the stream function from
 the third to the eighth snapshot, every two snapshots.
 
 ::
 
-   % stagpy field -o T -s 1:5 --vmin=0.8 --vmax=1.0
+    % stagpy field -o T -s 1,5 --vmin=0.8 --vmax=1.0
    
-will plot the temperature field from the first to the fifth snapshot while keeping the range of the colorbar fixed between 0.8 and 1. 
+will plot the temperature field from the first and the fifth snapshot while
+keeping the range of the colorbar fixed between 0.8 and 1.
+
+List of variables to plot
+-------------------------
+
+The list of fields specified with the ``-o`` (or ``--plot``) option follows the
+same rules as for the ``rprof`` and ``time`` subcommands.  Namely,
+``,``-separated variables are on the same subplots; ``.``-separated variables
+are on the same figure but different subplots; ``-``-separated variables are
+on different figures.
+
+Note that only two fields can be on the same subplot, the first field is a
+color map and the second field can be either:
+
+- a scalar field, isocontours are added to the plot;
+- a vector field (e.g. ``v`` for the ``(v1, v2, v3)`` vector), arrows are added
+  to the plot.
+
+For example, ``-o=T,v3`` asks for a temperature map with isocontour of the
+vertical velocity, while ``-o=T,v`` asks for a temperature map with velocity
+vectors on top of it.
+
+If you ask for more than two fields on the same subplot, extra fields are
+ignored.  ``-o=T,stream,v`` is therefore equivalent to ``-o=T,stream``.

docs/sources/cookbook/time.rst

@@ -33,7 +33,7 @@ time 0.03.
 
 ::
 
-    % stagpy time -o vrms_Tmin,Tmean,Tmax.dTdt
+    % stagpy time -o vrms-Tmin,Tmean,Tmax.dTdt
 
 creates two figures. The first one contains the time series of the rms
 velocity. The second one contains two subplots, the first one with the time
@@ -42,7 +42,7 @@ time derivative of the mean temperature. The variable names can be found by
 running the ``% stagpy var`` command. The variables you want on the same
 subplot are separated by commas ``,``, the variables you want on different
 subplots are separated by dots ``.``, and the variables you want on different
-figures are separated by underscores ``_``.
+figures are separated by dashes ``-``.
 
 ::
 

docs/sources/tuto.rst

@@ -35,7 +35,7 @@ to have some help on the available options for one particular sub command.
 
 A simple example would be::
 
-    % stagpy field -p path/to/run/ -o T,p -s 42
+    % stagpy field -p path/to/run/ -o T-p -s 42
 
 This asks StagPy to plot the temperature and pressure fields of snapshot 42
 of the run lying in ``./path/to/run``. When not specified, the path defaults to

stagpy/config.py

@@ -96,7 +96,7 @@ def _index_collection(arg):
 
 CONF_DEF['field'] = OrderedDict((
     ('plot',
-        Conf('T+stream', True, 'o',
+        Conf('T,stream', True, 'o',
              {'nargs': '?', 'const': '', 'type': str},
              True, 'variables to plot (see stagpy var)')),
     ('perturbation', switch_opt(False, None,
@@ -105,9 +105,6 @@ def _index_collection(arg):
                    False, 'Shift plot horizontally')),
     ('interpolate', switch_opt(True, None, 'apply Gouraud shading')),
     ('colorbar', switch_opt(True, None, 'add color bar to plot')),
-    ('shrinkcb',
-        Conf(0.5, False, None, {},
-             True, 'color bar shrink factor')),
     ('ix', Conf(None, True, None, {'type': int},
                 False, 'x-index of slice for 3D fields')),
     ('iy', Conf(None, True, None, {'type': int},

stagpy/field.py

@@ -1,9 +1,12 @@
 """Plot scalar and vector fields."""
 
+from itertools import chain
+
 import numpy as np
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import matplotlib.patches as mpat
+from mpl_toolkits.axes_grid1 import make_axes_locatable
 
 from . import conf, misc, phyvars
 from .error import NotAvailableError
@@ -119,28 +122,6 @@ def get_meshes_vec(step, var):
     return xmesh, ymesh, vec1, vec2
 
 
-def set_of_vars(arg_plot):
-    """Build set of needed field variables.
-
-    Each var is a tuple, first component is a scalar field, second component is
-    either:
-
-    - a scalar field, isocontours are added to the plot.
-    - a vector field (e.g. 'v' for the (v1,v2,v3) vector), arrows are added to
-      the plot.
-
-    Args:
-        arg_plot (str): string with variable names separated with
-            ``,`` (figures), and ``+`` (same plot).
-    Returns:
-        set of str: set of needed field variables.
-    """
-    sovs = set(tuple((var + '+').split('+')[:2])
-               for var in arg_plot.split(','))
-    sovs.discard(('', ''))
-    return sovs
-
-
 def plot_scalar(step, var, field=None, axis=None, **extra):
     """Plot scalar field.
 
@@ -217,13 +198,15 @@ def plot_scalar(step, var, field=None, axis=None, **extra):
 
     cbar = None
     if conf.field.colorbar:
-        cbar = plt.colorbar(surf, shrink=conf.field.shrinkcb)
+        cax = make_axes_locatable(axis).append_axes(
+            'right', size="3%", pad=0.15)
+        cbar = plt.colorbar(surf, cax=cax)
         cbar.set_label(meta.description +
                        (' pert.' if conf.field.perturbation else '') +
                        (' ({})'.format(unit) if unit else ''))
     if step.geom.spherical or conf.plot.ratio is None:
-        plt.axis('equal')
-        plt.axis('off')
+        axis.set_aspect('equal')
+        axis.set_axis_off()
     else:
         axis.set_aspect(conf.plot.ratio / axis.get_data_ratio())
     axis.set_adjustable('box')
@@ -283,13 +266,16 @@ def cmd():
         conf.core
     """
     sdat = StagyyData()
-    sovs = set_of_vars(conf.field.plot)
+    lovs = misc.list_of_vars(conf.field.plot)
+    # no more than two fields in a subplot
+    lovs = [[slov[:2] for slov in plov] for plov in lovs]
     minmax = {}
     if conf.plot.cminmax:
         conf.plot.vmin = None
         conf.plot.vmax = None
+        sovs = set(slov[0] for plov in lovs for slov in plov)
         for step in sdat.walk.filter(snap=True):
-            for var, _ in sovs:
+            for var in sovs:
                 if var in step.fields:
                     if var in phyvars.FIELD:
                         dim = phyvars.FIELD[var].dim
@@ -302,18 +288,23 @@ def cmd():
                     else:
                         minmax[var] = np.nanmin(field), np.nanmax(field)
     for step in sdat.walk.filter(snap=True):
-        for var in sovs:
-            if var[0] not in step.fields:
-                print("'{}' field on snap {} not found".format(var[0],
-                                                               step.isnap))
-                continue
-            opts = {}
-            if var[0] in minmax:
-                opts = dict(vmin=minmax[var[0]][0], vmax=minmax[var[0]][1])
-            fig, axis, _, _ = plot_scalar(step, var[0], **opts)
-            if valid_field_var(var[1]):
-                plot_iso(axis, step, var[1])
-            elif var[1]:
-                plot_vec(axis, step, var[1])
-            oname = '{}_{}'.format(*var) if var[1] else var[0]
+        for vfig in lovs:
+            fig, axes = plt.subplots(ncols=len(vfig), squeeze=False,
+                                     figsize=(12 * len(vfig), 9))
+            for axis, var in zip(axes[0], vfig):
+                if var[0] not in step.fields:
+                    print("'{}' field on snap {} not found".format(var[0],
+                                                                   step.isnap))
+                    continue
+                opts = {}
+                if var[0] in minmax:
+                    opts = dict(vmin=minmax[var[0]][0], vmax=minmax[var[0]][1])
+                plot_scalar(step, var[0], axis=axis, **opts)
+                if len(var) == 2:
+                    if valid_field_var(var[1]):
+                        plot_iso(axis, step, var[1])
+                    elif valid_field_var(var[1] + '1'):
+                        plot_vec(axis, step, var[1])
+            oname = '_'.join(chain.from_iterable(vfig))
+            plt.tight_layout(w_pad=3)
             misc.saveplot(fig, oname, step.isnap)

stagpy/misc.py

@@ -75,7 +75,7 @@ def list_of_vars(arg_plot):
 
     Args:
         arg_plot (str): string with variable names separated with
-            ``_`` (figures), ``.`` (subplots) and ``,`` (same subplot).
+            ``-`` (figures), ``.`` (subplots) and ``,`` (same subplot).
     Returns:
         three nested lists of str
 

tests/test_cli.py

@@ -10,7 +10,8 @@ def dir_isnap(request, repo_dir):
 
 @pytest.fixture(params=[
     ('stagpy field', ['stagpy_T_stream{:05d}.pdf']),
-    ('stagpy field -o=T,v3', ['stagpy_T{:05d}.pdf',
+    ('stagpy field -o=T.v3', ['stagpy_T_v3{:05d}.pdf']),
+    ('stagpy field -o=T-v3', ['stagpy_T{:05d}.pdf',
                               'stagpy_v3{:05d}.pdf']),
 ])
 def all_cmd_field(request, dir_isnap):

tests/test_field.py

@@ -23,9 +23,3 @@ def test_get_meshes_vec(step):
     xmesh, ymesh, vec1, vec2 = stagpy.field.get_meshes_vec(step, 'v')
     assert len(vec1.shape) == 2
     assert xmesh.shape == ymesh.shape == vec1.shape == vec2.shape
-
-
-def test_set_of_vars():
-    sovs = stagpy.field.set_of_vars('a+b,c,++,,d++e')
-    expected = set([('a', 'b'), ('c', ''), ('d', '')])
-    assert sovs == expected