processing.*_prof return Rprof instances

Similar to stream_function returning a Field.  Varr and Rprof are
moved to the datatypes module.

Author: amorison

docs/sources/apiref/datatypes.rst

@@ -5,3 +5,5 @@ datatypes
 
    .. autoclass:: Varf
    .. autoclass:: Field
+   .. autoclass:: Varr
+   .. autoclass:: Rprof

docs/sources/apiref/phyvars.rst

@@ -21,8 +21,6 @@ phyvars
       Dictionary of surface scalar fields output by StagYY. Keys are the
       variable names, values are :class:`Varf` instances.
 
-   .. autoclass:: Varr
-
    .. data:: RPROF
       :annotation: = {rprofvar: Varr()}
 

docs/sources/apiref/step.rst

@@ -5,6 +5,4 @@ _step
    :members:
    :private-members:
    :exclude-members: _init_shape, _get_raw_data, _scale_radius_mo,
-                     Rprof, _present_fields
-
-   .. autoclass:: Rprof
+                     _present_fields

stagpy/_step.py

@@ -9,21 +9,20 @@
 from __future__ import annotations
 from collections import abc
 from itertools import chain
-from typing import NamedTuple, TYPE_CHECKING
+import typing
 
 import numpy as np
 
-from . import error, phyvars, stagyyparsers, _helpers
+from . import error, phyvars, stagyyparsers
 from ._helpers import CachedReadOnlyProperty as crop
-from .datatypes import Field
+from .datatypes import Field, Rprof, Varr
 
-if TYPE_CHECKING:
+if typing.TYPE_CHECKING:
     from typing import (Dict, Any, Mapping, List, Iterator, Tuple, Optional,
                         Callable)
     from numpy import ndarray, signedinteger
     from pandas import DataFrame, Series
     from .datatypes import Varf
-    from .phyvars import Varr
     from .stagyydata import StagyyData
 
 
@@ -407,30 +406,14 @@ def __iter__(self):
         raise TypeError('tracers collection is not iterable')
 
 
-class Rprof(NamedTuple):
-    """Radial profile with associated radius and metadata.
-
-    Attributes:
-        values: the profile itself.
-        rad: the radial position.
-        meta: the metadata of the profile.
-    """
-
-    values: ndarray
-    rad: ndarray
-    meta: Varr
-
-
 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
+    defined in :data:`stagpy.phyvars.RPROF[_EXTRA]`.  Items are
+    :class:`stagpy.datatypes.Rprof` instances.  Note that
     profiles are automatically scaled if conf.scaling.dimensional is True.
 
     Attributes:
@@ -462,15 +445,12 @@ def __getitem__(self, name: str) -> Rprof:
             if name in phyvars.RPROF:
                 meta = phyvars.RPROF[name]
             else:
-                meta = phyvars.Varr(name, '', '1')
+                meta = Varr(name, '', '1')
         elif name in self._cached_extra:
             rprof, rad, meta = self._cached_extra[name]
         elif name in phyvars.RPROF_EXTRA:
-            meta = phyvars.RPROF_EXTRA[name]
-            rprof, rad = meta.description(step)
-            meta = phyvars.Varr(_helpers.baredoc(meta.description),
-                                meta.kind, meta.dim)
-            self._cached_extra[name] = Rprof(rprof, rad, meta)
+            self._cached_extra[name] = phyvars.RPROF_EXTRA[name](step)
+            rprof, rad, meta = self._cached_extra[name]
         else:
             raise error.UnknownRprofVarError(name)
         rprof, _ = step.sdat.scale(rprof, meta.dim)

stagpy/datatypes.py

@@ -30,3 +30,35 @@ class Field(NamedTuple):
 
     values: ndarray
     meta: Varf
+
+
+class Varr(NamedTuple):
+    """Metadata of radial profiles.
+
+    Attributes:
+        description: short description of the variable if it is output by
+            StagYY, function to compute it otherwise.
+        kind: shorter description to group similar variables under the same
+            label.
+        dim: dimension used to :func:`~stagpy.stagyydata.StagyyData.scale` to
+            dimensional values.
+    """
+
+    description: str
+    # Callable[[Step], Tuple[ndarray, ndarray]]]
+    kind: str
+    dim: str
+
+
+class Rprof(NamedTuple):
+    """Radial profile with associated radius and metadata.
+
+    Attributes:
+        values: the profile itself.
+        rad: the radial position.
+        meta: the metadata of the profile.
+    """
+
+    values: ndarray
+    rad: ndarray
+    meta: Varr

stagpy/phyvars.py

@@ -11,12 +11,11 @@
 from typing import NamedTuple, TYPE_CHECKING
 
 from . import processing
-from .datatypes import Varf
+from .datatypes import Varf, Varr
 
 if TYPE_CHECKING:
     from typing import Union, Callable, Tuple
     from numpy import ndarray
-    from ._step import Step
     from .stagyydata import StagyyData
 
 
@@ -138,23 +137,6 @@
 })
 
 
-class Varr(NamedTuple):
-    """Metadata of radial profiles.
-
-    Attributes:
-        description: short description of the variable if it is output by
-            StagYY, function to compute it otherwise.
-        kind: shorter description to group similar variables under the same
-            label.
-        dim: dimension used to :func:`~stagpy.stagyydata.StagyyData.scale` to
-            dimensional values.
-    """
-
-    description: Union[str, Callable[[Step], Tuple[ndarray, ndarray]]]
-    kind: str
-    dim: str
-
-
 RPROF = MappingProxyType({
     'r': Varr('Radial coordinate', 'Radius', 'm'),
     'Tmean': Varr('Temperature', 'Temperature', 'K'),
@@ -225,16 +207,16 @@ class Varr(NamedTuple):
 })
 
 RPROF_EXTRA = MappingProxyType({
-    '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'),
-    'advts': Varr(processing.advts_prof, 'Heat flux', 'W/m2'),
-    'advds': Varr(processing.advds_prof, 'Heat flux', 'W/m2'),
-    'advas': Varr(processing.advas_prof, 'Heat flux', 'W/m2'),
-    'energy': Varr(processing.energy_prof, 'Heat flux', 'W/m2'),
-    'ciover': Varr(processing.init_c_overturn, 'Concentration', '1'),
-    'cfover': Varr(processing.c_overturned, 'Concentration', '1'),
-    'advth': Varr(processing.advth, 'Heat Flux', 'W/m2'),
+    'dr': processing.delta_r,
+    'diff': processing.diff_prof,
+    'diffs': processing.diffs_prof,
+    'advts': processing.advts_prof,
+    'advds': processing.advds_prof,
+    'advas': processing.advas_prof,
+    'energy': processing.energy_prof,
+    'ciover': processing.init_c_overturn,
+    'cfover': processing.c_overturned,
+    'advth': processing.advth,
 })
 
 

stagpy/processing.py

@@ -12,7 +12,7 @@
 from scipy import integrate
 
 from .error import NotAvailableError
-from .datatypes import Field, Varf
+from .datatypes import Field, Varf, Rprof, Varr
 
 if typing.TYPE_CHECKING:
     from typing import Tuple
@@ -94,16 +94,17 @@ def mobility(sdat: StagyyData) -> Tuple[ndarray, ndarray]:
     return np.array(mob), np.array(time)
 
 
-def delta_r(step: Step) -> Tuple[ndarray, ndarray]:
-    """Cells thickness.
+def delta_r(step: Step) -> Rprof:
+    """Compute cell thicknesses.
 
     Args:
         step: a :class:`~stagpy._step.Step` of a StagyyData instance.
     Returns:
         the thickness of the cells and radius.
     """
     edges = step.rprofs.walls
-    return (edges[1:] - edges[:-1]), step.rprofs.centers
+    meta = Varr("Cell thickness", 'dr', 'm')
+    return Rprof((edges[1:] - edges[:-1]), step.rprofs.centers, meta)
 
 
 def _scale_prof(step: Step, rprof: ndarray, rad: ndarray = None) -> ndarray:
@@ -116,8 +117,8 @@ def _scale_prof(step: Step, rprof: ndarray, rad: ndarray = None) -> ndarray:
     return rprof * (2 * rad / (rtop + rbot))**2
 
 
-def diff_prof(step: Step) -> Tuple[ndarray, ndarray]:
-    """Diffusion.
+def diff_prof(step: Step) -> Rprof:
+    """Compute diffusion flux.
 
     Args:
         step: a :class:`~stagpy._step.Step` of a StagyyData instance.
@@ -132,11 +133,12 @@ def diff_prof(step: Step) -> Tuple[ndarray, ndarray]:
     diff = np.insert(diff, 0, (1 - tprof[0]) / (rad[0] - rbot))
     # assume ttop = 0
     diff = np.append(diff, tprof[-1] / (rtop - rad[-1]))
-    return diff, step.rprofs.walls
+    meta = Varr("Diffusion", 'Heat flux', 'W/m2')
+    return Rprof(diff, step.rprofs.walls, meta)
 
 
-def diffs_prof(step: Step) -> Tuple[ndarray, ndarray]:
-    """Scaled diffusion.
+def diffs_prof(step: Step) -> Rprof:
+    """Compute scaled diffusion flux.
 
     This computation takes sphericity into account if necessary.
 
@@ -145,12 +147,13 @@ def diffs_prof(step: Step) -> Tuple[ndarray, ndarray]:
     Returns:
         the diffusion and radius.
     """
-    diff, rad = diff_prof(step)
-    return _scale_prof(step, diff, rad), rad
+    diff, rad, _ = diff_prof(step)
+    meta = Varr("Scaled diffusion", 'Heat flux', 'W/m2')
+    return Rprof(_scale_prof(step, diff, rad), rad, meta)
 
 
-def advts_prof(step: Step) -> Tuple[ndarray, ndarray]:
-    """Scaled advection.
+def advts_prof(step: Step) -> Rprof:
+    """Compute scaled advection flux.
 
     This computation takes sphericity into account if necessary.
 
@@ -159,11 +162,13 @@ def advts_prof(step: Step) -> Tuple[ndarray, ndarray]:
     Returns:
         the scaled advection and radius.
     """
-    return _scale_prof(step, step.rprofs['advtot'].values), step.rprofs.centers
+    return Rprof(_scale_prof(step, step.rprofs['advtot'].values),
+                 step.rprofs.centers,
+                 Varr("Scaled advection", 'Heat flux', 'W/m2'))
 
 
-def advds_prof(step: Step) -> Tuple[ndarray, ndarray]:
-    """Scaled downward advection.
+def advds_prof(step: Step) -> Rprof:
+    """Compute scaled downward advection flux.
 
     This computation takes sphericity into account if necessary.
 
@@ -172,12 +177,13 @@ def advds_prof(step: Step) -> Tuple[ndarray, ndarray]:
     Returns:
         the scaled downward advection and radius.
     """
-    return (_scale_prof(step, step.rprofs['advdesc'].values),
-            step.rprofs.centers)
+    return Rprof(_scale_prof(step, step.rprofs['advdesc'].values),
+                 step.rprofs.centers,
+                 Varr("Scaled downward advection", 'Heat flux', 'W/m2'))
 
 
-def advas_prof(step: Step) -> Tuple[ndarray, ndarray]:
-    """Scaled upward advection.
+def advas_prof(step: Step) -> Rprof:
+    """Compute scaled upward advection flux.
 
     This computation takes sphericity into account if necessary.
 
@@ -186,11 +192,13 @@ def advas_prof(step: Step) -> Tuple[ndarray, ndarray]:
     Returns:
         the scaled upward advection and radius.
     """
-    return _scale_prof(step, step.rprofs['advasc'].values), step.rprofs.centers
+    return Rprof(_scale_prof(step, step.rprofs['advasc'].values),
+                 step.rprofs.centers,
+                 Varr("Scaled upward advection", 'Heat flux', 'W/m2'))
 
 
-def energy_prof(step: Step) -> Tuple[ndarray, ndarray]:
-    """Energy flux.
+def energy_prof(step: Step) -> Rprof:
+    """Compute total heat flux.
 
     This computation takes sphericity into account if necessary.
 
@@ -199,13 +207,14 @@ def energy_prof(step: Step) -> Tuple[ndarray, ndarray]:
     Returns:
         the energy flux and radius.
     """
-    diff, rad = diffs_prof(step)
-    adv, _ = advts_prof(step)
-    return (diff + np.append(adv, 0)), rad
+    diff, rad, _ = diffs_prof(step)
+    adv, _, _ = advts_prof(step)
+    return Rprof(diff + np.append(adv, 0), rad,
+                 Varr("Total heat flux", 'Heat flux', 'W/m2'))
 
 
-def advth(step: Step) -> Tuple[ndarray, ndarray]:
-    """Theoretical advection.
+def advth(step: Step) -> Rprof:
+    """Compute theoretical advection.
 
     This compute the theoretical profile of total advection as function of
     radius.
@@ -225,11 +234,12 @@ def advth(step: Step) -> Tuple[ndarray, ndarray]:
         th_adv = rad - rtop
     th_adv *= radio
     th_adv += step.timeinfo['Nutop']
-    return th_adv, rad
+    return Rprof(th_adv, rad,
+                 Varr("Theoretical advection", 'Heat flux', 'W/m2'))
 
 
-def init_c_overturn(step: Step) -> Tuple[ndarray, ndarray]:
-    """Initial concentration.
+def init_c_overturn(step: Step) -> Rprof:
+    """Compute concentration before overturn.
 
     This compute the resulting composition profile if fractional
     crystallization of a SMO is assumed.
@@ -257,11 +267,12 @@ def initprof(rpos):
 
     rad = np.linspace(rbot, rtop, 500)
     initprof = np.vectorize(initprof)
-    return initprof(rad), rad
+    return Rprof(initprof(rad), rad,
+                 Varr("Concentration before overturn", 'Concentration', '1'))
 
 
-def c_overturned(step: Step) -> Tuple[ndarray, ndarray]:
-    """Theoretical overturned concentration.
+def c_overturned(step: Step) -> Rprof:
+    """Compute theoretical overturned concentration.
 
     This compute the resulting composition profile if fractional
     crystallization of a SMO is assumed and then a purely radial
@@ -273,9 +284,10 @@ def c_overturned(step: Step) -> Tuple[ndarray, ndarray]:
         the composition and radius.
     """
     rbot, rtop = step.rprofs.bounds
-    cinit, rad = init_c_overturn(step)
-    radf = (rtop**3 + rbot**3 - rad**3)**(1 / 3)
-    return cinit, radf
+    cinit = init_c_overturn(step)
+    radf = (rtop**3 + rbot**3 - cinit.rad**3)**(1 / 3)
+    return Rprof(cinit.values, radf,
+                 Varr("Overturned concentration", 'Concentration', '1'))
 
 
 def stream_function(step: Step) -> Field: