Time processing function return Tseries instances

Similar to what was done with Rprof and Field.  Tseries and Vart are
now in datatypes.

Author: amorison

docs/sources/apiref/datatypes.rst

@@ -7,3 +7,5 @@ datatypes
    .. autoclass:: Field
    .. autoclass:: Varr
    .. autoclass:: Rprof
+   .. autoclass:: Vart
+   .. autoclass:: Tseries

docs/sources/apiref/phyvars.rst

@@ -39,8 +39,6 @@ phyvars
       Dictionary of radial profiles of the reference state. Keys are the
       variable names, values are :class:`Varr` instances.
 
-   .. autoclass:: Vart
-
    .. data:: TIME
       :annotation: = {timevar: Vart()}
 

stagpy/commands.py

@@ -50,8 +50,8 @@ def info_cmd():
             series = series.copy()
             dimensions = []
             for var, val in series.iteritems():
-                dim = phyvars.TIME.get(var) or phyvars.TIME_EXTRA.get(var)
-                dim = dim.dim if dim is not None else '1'
+                meta = phyvars.TIME.get(var)
+                dim = meta.dim if meta is not None else '1'
                 if dim == '1':
                     dimensions.append('')
                 else:

stagpy/datatypes.py

@@ -62,3 +62,34 @@ class Rprof(NamedTuple):
     values: ndarray
     rad: ndarray
     meta: Varr
+
+
+class Vart(NamedTuple):
+    """Metadata of time series.
+
+    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
+    kind: str
+    dim: str
+
+
+class Tseries(NamedTuple):
+    """A time series with associated time and metadata.
+
+    Attributes:
+        values: the series itself.
+        time: the time vector.
+        meta: the metadata of the series.
+    """
+
+    values: ndarray
+    time: ndarray
+    meta: Vart

stagpy/phyvars.py

@@ -5,18 +5,11 @@
 be computed from other variables.
 """
 
-from __future__ import annotations
 from operator import attrgetter
 from types import MappingProxyType
-from typing import NamedTuple, TYPE_CHECKING
 
 from . import processing
-from .datatypes import Varf, Varr
-
-if TYPE_CHECKING:
-    from typing import Union, Callable, Tuple
-    from numpy import ndarray
-    from .stagyydata import StagyyData
+from .datatypes import Varf, Varr, Vart
 
 
 FIELD = MappingProxyType({
@@ -220,23 +213,6 @@
 })
 
 
-class Vart(NamedTuple):
-    """Metadata of time series.
-
-    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[[StagyyData], Tuple[ndarray, ndarray]]]
-    kind: str
-    dim: str
-
-
 TIME = MappingProxyType({
     't': Vart('Time', 'Time', 's'),
     'ftop': Vart('Heat flux at top', 'Heat flux', 'W/m2'),
@@ -270,10 +246,10 @@ class Vart(NamedTuple):
 })
 
 TIME_EXTRA = MappingProxyType({
-    'dt': Vart(processing.dtime, 'dt', 's'),
-    'dTdt': Vart(processing.dt_dt, r'dT/dt', 'K/s'),
-    'ebalance': Vart(processing.ebalance, r'$\mathrm{Nu}$', '1'),
-    'mobility': Vart(processing.mobility, 'Mobility', '1'),
+    'dt': processing.dtime,
+    'dTdt': processing.dt_dt,
+    'ebalance': processing.ebalance,
+    'mobility': processing.mobility,
 })
 
 REFSTATE = MappingProxyType({

stagpy/processing.py

@@ -12,17 +12,16 @@
 from scipy import integrate
 
 from .error import NotAvailableError
-from .datatypes import Field, Varf, Rprof, Varr
+from .datatypes import Field, Varf, Rprof, Varr, Tseries, Vart
 
 if typing.TYPE_CHECKING:
-    from typing import Tuple
     from numpy import ndarray
     from .stagyydata import StagyyData
     from ._step import Step
 
 
-def dtime(sdat: StagyyData) -> Tuple[ndarray, ndarray]:
-    """Time increment dt.
+def dtime(sdat: StagyyData) -> Tseries:
+    """Compute time increment.
 
     Compute dt as a function of time.
 
@@ -32,11 +31,12 @@ def dtime(sdat: StagyyData) -> Tuple[ndarray, ndarray]:
         dt and time arrays.
     """
     time = sdat.tseries.time
-    return time[1:] - time[:-1], time[:-1]
+    return Tseries(time[1:] - time[:-1], time[:-1],
+                   Vart("Time increment dt", 'dt', 's'))
 
 
-def dt_dt(sdat: StagyyData) -> Tuple[ndarray, ndarray]:
-    """Derivative of temperature.
+def dt_dt(sdat: StagyyData) -> Tseries:
+    """Compute derivative of temperature.
 
     Compute dT/dt as a function of time using an explicit Euler scheme.
 
@@ -47,11 +47,12 @@ def dt_dt(sdat: StagyyData) -> Tuple[ndarray, ndarray]:
     """
     temp, time, _ = sdat.tseries['Tmean']
     dtdt = (temp[1:] - temp[:-1]) / (time[1:] - time[:-1])
-    return dtdt, time[:-1]
+    return Tseries(dtdt, time[:-1],
+                   Vart("Derivative of temperature", r'dT/dt', 'K/s'))
 
 
-def ebalance(sdat: StagyyData) -> Tuple[ndarray, ndarray]:
-    """Energy balance.
+def ebalance(sdat: StagyyData) -> Tseries:
+    """Compute energy balance.
 
     Compute Nu_t - Nu_b + V*dT/dt as a function of time using an explicit
     Euler scheme. This should be zero if energy is conserved.
@@ -68,16 +69,17 @@ def ebalance(sdat: StagyyData) -> Tuple[ndarray, ndarray]:
     else:
         coefsurf = 1.
         volume = 1.
-    dtdt, time = dt_dt(sdat)
+    dtdt, time, _ = dt_dt(sdat)
     ftop = sdat.tseries['ftop'].values * coefsurf
     fbot = sdat.tseries['fbot'].values
     radio = sdat.tseries['H_int'].values
     ebal = ftop[1:] - fbot[1:] + volume * (dtdt - radio[1:])
-    return ebal, time
+    return Tseries(ebal, time,
+                   Vart("Energy balance", r'$\mathrm{Nu}$', '1'))
 
 
-def mobility(sdat: StagyyData) -> Tuple[ndarray, ndarray]:
-    """Plates mobility.
+def mobility(sdat: StagyyData) -> Tseries:
+    """Compute plates mobility.
 
     Compute the ratio vsurf / vrms.
 
@@ -91,7 +93,8 @@ def mobility(sdat: StagyyData) -> Tuple[ndarray, ndarray]:
     for step in sdat.steps.filter(rprofs=True):
         time.append(step.timeinfo['t'])
         mob.append(step.rprofs['vrms'].values[-1] / step.timeinfo['vrms'])
-    return np.array(mob), np.array(time)
+    return Tseries(np.array(mob), np.array(time),
+                   Vart("Plates mobility", 'Mobility', '1'))
 
 
 def delta_r(step: Step) -> Rprof:

stagpy/stagyydata.py

@@ -9,14 +9,13 @@
 
 import re
 import pathlib
-from collections import namedtuple
 from itertools import zip_longest
 
 import numpy as np
 
 from . import conf, error, parfile, phyvars, stagyyparsers, _helpers, _step
 from ._helpers import CachedReadOnlyProperty as crop
-from .datatypes import Rprof
+from .datatypes import Rprof, Tseries, Vart
 
 
 def _as_view_item(obj):
@@ -159,20 +158,15 @@ def adiabats(self):
         return self._data[1]
 
 
-Tseries = namedtuple('Tseries', ['values', 'time', 'meta'])
-
-
 class _Tseries:
     """Time series.
 
     The :attr:`StagyyData.tseries` attribute is an instance of this class.
 
     :class:`_Tseries` implements the getitem mechanism.  Keys are series names
-    defined in :data:`stagpy.phyvars.TIME[_EXTRA]`.  An item is a named tuple
-    ('values', 'time', 'meta'), respectively the series itself, the time at
-    which it is evaluated, and meta is a :class:`stagpy.phyvars.Vart` instance
-    with relevant metadata.  Note that series are automatically scaled if
-    conf.scaling.dimensional is True.
+    defined in :data:`stagpy.phyvars.TIME[_EXTRA]`.  Items are
+    :class:`stagpy.datatypes.Tseries` instances.  Note that series are
+    automatically scaled if conf.scaling.dimensional is True.
 
     Attributes:
         sdat (:class:`StagyyData`): the StagyyData instance owning the
@@ -210,15 +204,12 @@ def __getitem__(self, name):
             if name in phyvars.TIME:
                 meta = phyvars.TIME[name]
             else:
-                meta = phyvars.Vart(name, None, '1')
+                meta = Vart(name, '', '1')
         elif name in self._cached_extra:
             series, time, meta = self._cached_extra[name]
         elif name in phyvars.TIME_EXTRA:
-            meta = phyvars.TIME_EXTRA[name]
-            series, time = meta.description(self.sdat)
-            meta = phyvars.Vart(_helpers.baredoc(meta.description),
-                                meta.kind, meta.dim)
-            self._cached_extra[name] = series, time, meta
+            self._cached_extra[name] = phyvars.TIME_EXTRA[name](self.sdat)
+            series, time, meta = self._cached_extra[name]
         else:
             raise error.UnknownTimeVarError(name)
         series, _ = self.sdat.scale(series, meta.dim)

tests/test_phyvars.py

@@ -4,7 +4,7 @@
 
 def test_dim():
     allvars = chain(phyvars.FIELD.values(), phyvars.RPROF.values(),
-                    phyvars.TIME.values(), phyvars.TIME_EXTRA.values())
+                    phyvars.TIME.values())
     for var in allvars:
         if var.dim != '1':
             assert var.dim in phyvars.SCALES

tests/test_processing.py

@@ -1,19 +1,22 @@
 from stagpy import processing, phyvars
 
 
+def tseries_checks(tseries, expected_size):
+    assert tseries.values.shape == tseries.time.shape == (expected_size,)
+    assert tseries.meta.dim == '1' or tseries.meta.dim in phyvars.SCALES
+
+
 def rprof_checks(rprof, expected_size):
     assert rprof.values.shape == rprof.rad.shape == (expected_size,)
     assert rprof.meta.dim == '1' or rprof.meta.dim in phyvars.SCALES
 
 
 def test_dt_dt(sdat):
-    dtdt, time = processing.dt_dt(sdat)
-    assert dtdt.shape == time.shape == (sdat.tseries.time.shape[0] - 1,)
+    tseries_checks(processing.dt_dt(sdat), sdat.tseries.time.shape[0] - 1)
 
 
 def test_ebalance(sdat):
-    ebal, time = processing.ebalance(sdat)
-    assert ebal.shape == time.shape == (sdat.tseries.time.shape[0] - 1,)
+    tseries_checks(processing.ebalance(sdat), sdat.tseries.time.shape[0] - 1)
 
 
 def test_r_edges(step):