Merge branch 'rprof' into master

Author: amorison

docs/sources/stagyydata.rst

@@ -86,32 +86,24 @@ are given a default value according to the par file ``~/.config/stagpy/par``.
 Radial profiles
 ---------------
 
-Radial profile data are contained in the
-:attr:`~stagpy.stagyydata.StagyyData.rprof` attribute of a
-:class:`~stagpy.stagyydata.StagyyData` 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 a 2
-levels multi-index, the first level being the time step number (:attr:`istep`),
-and the second level being the cells 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:
 
-The radial profile of a given time step can be accessed from :attr:`_Step.rprof
-<stagpy.stagyydata._Step.rprof>`. For example, ``sdat.steps[1000].rprof`` is
-equivalent to ``sdat.rprof.loc[1000]``. The columns of the obtained dataframe
-are the variable names, and its index is the cells number.
+- :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:
 
-As an example, the following lines are two ways of accessing the horizontal
-average temperature in the bottom cell, at the 1000th timestep::
+    - :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.
 
-    # extract rprof data for the 1000th timestep,
-    # and then take the temperature in the bottom cell
-    sdat.rprof.loc[1000].loc[0,'Tmean']
-    # extract the temperature profile for the 1000th timestep,
-    # and then take the bottom cell
-    sdat.rprof.loc[1000,'Tmean'][0]
+The list of available variables can be obtained by running ``% stagpy var``.
 
-If the radial profiles of the 1000th timestep are not available, these would
-both result in a ``KeyError``.
+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.
-
-        Set to None if no radial profiles data is available for this time step.
-        """
-        if self.istep not in self.sdat.rprof.index.levels[0]:
-            return None
-        return self.sdat.rprof.loc[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,12 +593,11 @@ def main_plates(sdat):
     """Plot several plates information."""
     # calculating averaged horizontal surface velocity
     # needed for redimensionalisation
-    ilast = sdat.rprof.index.levels[0][-1]
-    rlast = sdat.rprof.loc[ilast]
+    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