New rprofs_averaged attribute of _StepsView

This gives access to time-average of radial profiles over the
_StepsView.  It has the same API as _Rprofs, with an extra steps
attribute pointing to the owner of the rprofs_averaged attribute.

Author: amorison

stagpy/_step.py

@@ -404,9 +404,10 @@ def _rprofs(self):
         if self._data is UNDETERMINED:
             step = self.step
             self._data = step.sdat._rprof_and_times[0].get(step.istep)
-            if self._data is None:
-                raise error.MissingDataError(
-                    f'No rprof data in step {step.istep} of {step.sdat}')
+        if self._data is None:
+            step = self.step
+            raise error.MissingDataError(
+                f'No rprof data in step {step.istep} of {step.sdat}')
         return self._data
 
     def __getitem__(self, name):

stagpy/rprof.py

@@ -1,10 +1,9 @@
 """Plot radial profiles."""
+import re
 
 import matplotlib.pyplot as plt
-import numpy as np
 
 from . import conf, misc
-from ._step import Rprof
 from .stagyydata import StagyyData
 
 
@@ -86,41 +85,21 @@ def plot_average(sdat, lovs):
         conf.core.snapshots: the slice of snapshots.
         conf.conf.timesteps: the slice of timesteps.
     """
-    steps_iter = iter(sdat.walk.filter(rprofs=True))
-    try:
-        step = next(steps_iter)
-    except StopIteration:
-        return
+    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)
 
-    istart = step.istep
-    nprofs = 1
-    rprof_averaged = {}
-    rads = {}
-    metas = {}
-
-    # assume constant z spacing for the moment
-    for rvar in sovs:
-        rprof, rads[rvar], metas[rvar] = step.rprofs[rvar]
-        rprof_averaged[rvar] = np.copy(rprof)
-
-    for step in steps_iter:
-        nprofs += 1
-        for rvar in sovs:
-            rprof_averaged[rvar] += step.rprofs[rvar].values
-
-    ilast = step.istep
-    for rvar in sovs:
-        # cast to float so that division happens in place
-        rprof_averaged[rvar] = Rprof(rprof_averaged[rvar] / nprofs,
-                                     rads[rvar], metas[rvar])
-    rcmb, rsurf = step.rprofs.bounds
+    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])
 
-    stepstr = f'{istart}_{ilast}'
-
     _plot_rprof_list(sdat, lovs, rprof_averaged, stepstr)
 
 

stagpy/stagyydata.py

@@ -12,6 +12,8 @@
 from collections import namedtuple
 from itertools import zip_longest
 
+import numpy as np
+
 from . import conf, error, misc, parfile, phyvars, stagyyparsers, _step
 from ._step import UNDETERMINED
 
@@ -259,6 +261,43 @@ def at_step(self, istep):
         return self._tseries.loc[istep]
 
 
+class _RprofsAveraged(_step._Rprofs):
+    """Radial profiles time-averaged over a :class:`_StepsView`.
+
+    The :attr:`_StepsView.rprofs_averaged` attribute is an instance of this
+    class.
+
+    It implements the same interface as :class:`~stagpy._step._Rprofs` but
+    returns time-averaged profiles instead.
+
+    Attributes:
+        steps (:class:`_StepsView`): the object owning the
+            :class:`_RprofsAveraged` instance
+    """
+
+    def __init__(self, steps):
+        self.steps = steps.filter(rprofs=True)
+        self._cached_data = {}
+        super().__init__(next(iter(self.steps)))
+
+    def __getitem__(self, name):
+        # the averaging method has two shortcomings:
+        # - does not take into account time changing geometry;
+        # - does not take into account time changing timestep.
+        if name in self._cached_data:
+            return self._cached_data[name]
+        steps_iter = iter(self.steps)
+        rprof, rad, meta = next(steps_iter).rprofs[name]
+        rprof = np.copy(rprof)
+        nprofs = 1
+        for step in steps_iter:
+            nprofs += 1
+            rprof += step.rprofs[name].values
+        rprof /= nprofs
+        self._cached_data[name] = _step.Rprof(rprof, rad, meta)
+        return self._cached_data[name]
+
+
 class _Steps:
     """Collections of time steps.
 
@@ -494,6 +533,7 @@ class _StepsView:
     def __init__(self, steps_col, items):
         self._col = steps_col
         self._items = items
+        self._rprofs_averaged = None
         self._flt = {
             'snap': False,
             'rprofs': False,
@@ -502,6 +542,12 @@ def __init__(self, steps_col, items):
         }
         self._dflt_func = self._flt['func']
 
+    @property
+    def rprofs_averaged(self):
+        if self._rprofs_averaged is None:
+            self._rprofs_averaged = _RprofsAveraged(self)
+        return self._rprofs_averaged
+
     def __repr__(self):
         rep = repr(self._col)
         items = []