New radial profiles interface

The previous interface was exposing a pandas dataframe, which itself
has an API that's unecessarily complex for our goals and not flexible
as subclassing dataframes properly is difficult.  In particular, this
led to the awkward rprof.get_rprof interface to get RPROF_EXTRA
variables.

This commit fixes these problems by implementing a rprof interface as
member of Step instances which abstracts rprof.get_rprof as well as
misc.get_rbounds.

In particular:

- step.rprofs replaces step.rprof (for consistency with tracers and
  fields attributes which have a plural name);
- items of rprofs are named tuple with values, rad, and meta fields;
- EXTRA variables are seamlessly accessible through the same interface;
- rprof.get_rprof is removed as it is no longer necessary;
- step.rprofs.bounds replaces misc.get_rbounds;
- step.rprofs.centers offers easy access to cell center location;
- step.rprofs.walls offers access to walls location and replaces
  the redges EXTRA variable (which is dropped).

Author: amorison

docs/sources/stagyydata.rst

@@ -86,14 +86,23 @@ are given a default value according to the par file ``~/.config/stagpy/par``.
 Radial profiles
 ---------------
 
-Radial profile data are contained in the :attr:`~stagpy._step.Step.rprof`
-attribute of a :class:`~stagpy._step.Step` instance.  This attribute is a
-:class:`pandas.DataFrame`. Its :attr:`columns` are the names of available
-variables (such as e.g. ``'Tmean'`` and ``'ftop'``).  Its :attr:`index` is the
-cell number (from ``0`` to ``nz-1``). The list of available variables can be
-obtained by running ``% stagpy var``.
+Radial profile data are accessible trough the :attr:`~stagpy._step.Step.rprofs`
+attribute of a :class:`~stagpy._step.Step` instance.  This attribute implements
+getitem to access radial profiles.  Keys are the names of available
+variables (such as e.g. ``'Tmean'`` and ``'ftop'``).  Items are named tuples
+with three fields:
+
+- :data:`values`: the profile itself;
+- :data:`rad`: the radial position at which the profile is evaluated;
+- :data:`meta`: metadata of the profile, also a named tuple with:
+
+    - :data:`description`: explanation of what the profile is;
+    - :data:`kind`: the category of profile;
+    - :data:`dim`: the dimension of the profile (if applicable) in SI units.
+
+The list of available variables can be obtained by running ``% stagpy var``.
 
-For example, ``sdat.steps[1000].rprof['Tmean']`` is the temperature profile of
+For example, ``sdat.steps[1000].rprofs['Tmean']`` is the temperature profile of
 the 1000th timestep.
 
 Time series

docs/sources/tuto.rst

@@ -87,7 +87,7 @@ Instantiating and using this class is rather simple::
     sdat = StagyyData('path/to/run/')
 
     # absolute vertical velocity profile of last snapshot
-    last_v_prof = sdat.snaps[-1].rprof['vzabs']
+    last_v_prof = sdat.snaps[-1].rprof['vzabs'].values
 
     # temperature field of the 10000th time step
     # (will be None if no snapshot is available at this timestep)

stagpy/_step.py

@@ -7,12 +7,13 @@
 """
 
 from collections.abc import Mapping
+from collections import namedtuple
 from itertools import chain
 import re
 
 import numpy as np
 
-from . import error, phyvars, stagyyparsers
+from . import error, misc, phyvars, stagyyparsers
 
 
 UNDETERMINED = object()
@@ -370,6 +371,103 @@ def __iter__(self):
         raise TypeError('tracers collection is not iterable')
 
 
+Rprof = namedtuple('Rprof', ['values', 'rad', 'meta'])
+
+
+class _Rprofs:
+    """Radial profiles data structure.
+
+    The :attr:`Step.rprofs` attribute is an instance of this class.
+
+    :class:`_Rprofs` implements the getitem mechanism.  Keys are profile names
+    defined in :data:`stagpy.phyvars.RPROF[_EXTRA]`.  An item is a named tuple
+    ('values', 'rad', 'meta'), respectively the profile itself, the radial
+    position at which it is evaluated, and meta is a
+    :class:`stagpy.phyvars.Varr` instance with relevant metadata.  Note that
+    profiles are automatically scaled if conf.scaling.dimensional is True.
+
+    Attributes:
+        step (:class:`Step`): the step object owning the :class:`_Rprofs`
+            instance
+    """
+
+    def __init__(self, step):
+        self.step = step
+        self._data = UNDETERMINED
+        self._centers = UNDETERMINED
+        self._walls = UNDETERMINED
+        self._bounds = UNDETERMINED
+
+    @property
+    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('No rprof data in step {} of {}'
+                                             .format(step.istep, step.sdat))
+        return self._data
+
+    def __getitem__(self, name):
+        step = self.step
+        if name in self._rprofs.columns:
+            rprof = self._rprofs[name].values
+            rad = self.centers
+            if name in phyvars.RPROF:
+                meta = phyvars.RPROF[name]
+            else:
+                meta = phyvars.Varr(name, None, '1')
+        elif name in phyvars.RPROF_EXTRA:
+            meta = phyvars.RPROF_EXTRA[name]
+            rprof, rad = meta.description(step)
+            meta = phyvars.Varr(misc.baredoc(meta.description),
+                                meta.kind, meta.dim)
+        else:
+            raise error.UnknownRprofVarError(name)
+        rprof, _ = step.sdat.scale(rprof, meta.dim)
+        rad, _ = step.sdat.scale(rad, 'm')
+
+        return Rprof(rprof, rad, meta)
+
+    @property
+    def centers(self):
+        """Radial position of cell centers."""
+        if self._centers is UNDETERMINED:
+            self._centers = self._rprofs['r'].values + self.bounds[0]
+        return self._centers
+
+    @property
+    def walls(self):
+        """Radial position of cell walls."""
+        if self._walls is UNDETERMINED:
+            rbot, rtop = self.bounds
+            centers = self.centers
+            # assume walls are mid-way between T-nodes
+            # could be T-nodes at center between walls
+            self._walls = (centers[:-1] + centers[1:]) / 2
+            self._walls = np.insert(self._walls, 0, rbot)
+            self._walls = np.append(self._walls, rtop)
+        return self._walls
+
+    @property
+    def bounds(self):
+        """Radial or vertical position of boundaries.
+
+        Radial/vertical positions of boundaries of the domain.
+        """
+        if self._bounds is UNDETERMINED:
+            step = self.step
+            if step.geom is not None:
+                rcmb = step.geom.rcmb
+            else:
+                rcmb = step.sdat.par['geometry']['r_cmb']
+                if step.sdat.par['geometry']['shape'].lower() == 'cartesian':
+                    rcmb = 0
+            rcmb = max(rcmb, 0)
+            self._bounds = rcmb, rcmb + 1
+        return self._bounds
+
+
 class Step:
     """Time step data structure.
 
@@ -416,6 +514,7 @@ def __init__(self, istep, sdat):
         self.sfields = _Fields(self, phyvars.SFIELD, [],
                                phyvars.SFIELD_FILES, phyvars.SFIELD_FILES_H5)
         self.tracers = _Tracers(self)
+        self.rprofs = _Rprofs(self)
         self._isnap = UNDETERMINED
 
     @property
@@ -438,16 +537,6 @@ def timeinfo(self):
             return None
         return self.sdat.tseries.loc[self.istep]
 
-    @property
-    def rprof(self):
-        """Radial profiles data of the time step.
-
-        This is a :class:`pandas.DataFrame` with iz as index and variable names
-        as columns.  Set to None if no radial profiles data is available for
-        this time step.
-        """
-        return self.sdat._rprof_and_times[0].get(self.istep)
-
     @property
     def isnap(self):
         """Snapshot index corresponding to time step.

stagpy/misc.py

@@ -121,26 +121,6 @@ def set_of_vars(lovs):
     return set(var for pvars in lovs for svars in pvars for var in svars)
 
 
-def get_rbounds(step):
-    """Radial or vertical position of boundaries.
-
-    Args:
-        step (:class:`~stagpy.stagyydata._Step`): a step of a StagyyData
-            instance.
-    Returns:
-        tuple of floats: radial or vertical positions of boundaries of the
-        domain.
-    """
-    if step.geom is not None:
-        rcmb = step.geom.rcmb
-    else:
-        rcmb = step.sdat.par['geometry']['r_cmb']
-        if step.sdat.par['geometry']['shape'].lower() == 'cartesian':
-            rcmb = 0
-    rcmb = max(rcmb, 0)
-    return rcmb, rcmb + 1
-
-
 class InchoateFiles:
     """Context manager handling files whose names are not known yet.
 

stagpy/phyvars.py

@@ -192,7 +192,6 @@
 ))
 
 RPROF_EXTRA = OrderedDict((
-    ('redges', Varr(processing.r_edges, 'Radius', 'm')),
     ('dr', Varr(processing.delta_r, 'dr', 'm')),
     ('diff', Varr(processing.diff_prof, 'Heat flux', 'W/m2')),
     ('diffs', Varr(processing.diffs_prof, 'Heat flux', 'W/m2')),

stagpy/plates.py

@@ -593,11 +593,11 @@ def main_plates(sdat):
     """Plot several plates information."""
     # calculating averaged horizontal surface velocity
     # needed for redimensionalisation
-    rlast = sdat.snaps[-1].rprof
+    rlast = sdat.snaps[-1].rprofs
     nprof = 0
-    uprof_averaged = rlast.loc[:, 'vhrms'] * 0
-    for step in sdat.walk.filter(rprof=True):
-        uprof_averaged += step.rprof['vhrms']
+    uprof_averaged = np.zeros_like(rlast['vhrms'].values)
+    for step in sdat.walk.filter(rprofs=True):
+        uprof_averaged += step.rprofs['vhrms'].values
         nprof += 1
     uprof_averaged /= nprof
     radius = rlast['r'].values