Use rprofs API to simplify rprof module

Author: amorison

stagpy/rprof.py

@@ -7,8 +7,24 @@
 from .stagyydata import StagyyData
 
 
-def _plot_rprof_list(sdat, lovs, rprofs, stepstr):
-    """Plot requested profiles."""
+def plot_rprofs(rprofs, lovs):
+    """Plot requested radial profiles.
+
+    Args:
+        rprofs (:class:`~stagpy._step._Rprofs`): a radial profile viewer
+            such as :attr:`Step.rprofs` and :attr:`_StepsView.rprofs_averaged`.
+        lovs (nested list of str): nested list of profile names such as the one
+            produced by :func:`stagpy.misc.list_of_vars`.
+    """
+    try:
+        stepstr = repr(rprofs.steps)
+        reg = re.compile(
+            r'^StagyyData\(.*\)\.(steps|snaps)\[(.*)\](?:.filter\(.*\))?$')
+        stepstr = '_'.join(reg.match(stepstr).groups())
+    except AttributeError:
+        stepstr = str(rprofs.step.istep)
+    sdat = rprofs.step.sdat
+
     for vfig in lovs:
         fig, axes = plt.subplots(ncols=len(vfig), sharey=True,
                                  figsize=(4 * len(vfig), 6))
@@ -20,12 +36,10 @@ def _plot_rprof_list(sdat, lovs, rprofs, stepstr):
             fname += '_'.join(vplt) + '_'
             for ivar, (rprof, rad, meta) in enumerate(profs_on_plt):
                 if conf.rprof.depth:
-                    rad = rprofs['bounds'][1] - rad
+                    rad = sdat.scale(rprofs.bounds[1], 'm')[0] - rad
                 axes[iplt].plot(rprof, rad,
                                 conf.rprof.style,
                                 label=meta.description)
-                if conf.rprof.depth:
-                    axes[iplt].invert_yaxis()
                 if xlabel is None:
                     xlabel = meta.kind
                 elif xlabel != meta.kind:
@@ -34,18 +48,18 @@ def _plot_rprof_list(sdat, lovs, rprofs, stepstr):
                 xlabel = meta.description
             if xlabel:
                 _, unit = sdat.scale(1, meta.dim)
-                if unit:
-                    xlabel += f' ({unit})'
+                xlabel += f' ({unit})' if unit else ''
                 axes[iplt].set_xlabel(xlabel)
             if vplt[0][:3] == 'eta':  # list of log variables
                 axes[iplt].set_xscale('log')
             axes[iplt].set_xlim(left=conf.plot.vmin, right=conf.plot.vmax)
             if ivar:
                 axes[iplt].legend()
+        if conf.rprof.depth:
+            axes[0].invert_yaxis()
         ylabel = 'Depth' if conf.rprof.depth else 'Radius'
         _, unit = sdat.scale(1, 'm')
-        if unit:
-            ylabel += f' ({unit})'
+        ylabel += f' ({unit})' if unit else ''
         axes[0].set_ylabel(ylabel)
         misc.saveplot(fig, fname + stepstr)
 
@@ -73,60 +87,6 @@ def plot_grid(step):
     misc.saveplot(fig, 'grid', step.istep)
 
 
-def plot_average(sdat, lovs):
-    """Plot time averaged profiles.
-
-    Args:
-        sdat (:class:`~stagpy.stagyydata.StagyyData`): a StagyyData instance.
-        lovs (nested list of str): nested list of profile names such as
-            the one produced by :func:`stagpy.misc.list_of_vars`.
-
-    Other Parameters:
-        conf.core.snapshots: the slice of snapshots.
-        conf.conf.timesteps: the slice of timesteps.
-    """
-    reg = re.compile(
-        r'^StagyyData\(.*\)\.(steps|snaps)\[(.*)\](?:.filter\(.*\))?$')
-    stepstr = repr(sdat.walk)
-    stepstr = '_'.join(reg.match(stepstr).groups())
-    rprofs = sdat.walk.rprofs_averaged
-
-    sovs = misc.set_of_vars(lovs)
-
-    rprof_averaged = {rvar: rprofs[rvar] for rvar in sovs}
-
-    rcmb, rsurf = rprofs.bounds
-    step = rprofs.step
-    rprof_averaged['bounds'] = (step.sdat.scale(rcmb, 'm')[0],
-                                step.sdat.scale(rsurf, 'm')[0])
-
-    _plot_rprof_list(sdat, lovs, rprof_averaged, stepstr)
-
-
-def plot_every_step(sdat, lovs):
-    """Plot profiles at each time step.
-
-    Args:
-        sdat (:class:`~stagpy.stagyydata.StagyyData`): a StagyyData instance.
-        lovs (nested list of str): nested list of profile names such as
-            the one produced by :func:`stagpy.misc.list_of_vars`.
-
-    Other Parameters:
-        conf.core.snapshots: the slice of snapshots.
-        conf.conf.timesteps: the slice of timesteps.
-    """
-    sovs = misc.set_of_vars(lovs)
-
-    for step in sdat.walk.filter(rprofs=True):
-        rprofs = {rvar: step.rprofs[rvar] for rvar in sovs}
-        rcmb, rsurf = step.rprofs.bounds
-        rprofs['bounds'] = (step.sdat.scale(rcmb, 'm')[0],
-                            step.sdat.scale(rsurf, 'm')[0])
-        stepstr = str(step.istep)
-
-        _plot_rprof_list(sdat, lovs, rprofs, stepstr)
-
-
 def cmd():
     """Implementation of rprof subcommand.
 
@@ -145,6 +105,7 @@ def cmd():
         return
 
     if conf.rprof.average:
-        plot_average(sdat, lovs)
+        plot_rprofs(sdat.walk.rprofs_averaged, lovs)
     else:
-        plot_every_step(sdat, lovs)
+        for step in sdat.walk.filter(rprofs=True):
+            plot_rprofs(step.rprofs, lovs)