Replace NamedTuples with dataclasses

Author: amorison

stagpy/_step.py

@@ -508,10 +508,16 @@ def __getitem__(self, name: str) -> Rprof:
             else:
                 meta = Varr(name, "", "1")
         elif name in self._cached_extra:
-            rprof, rad, meta = self._cached_extra[name]
+            rpf = self._cached_extra[name]
+            rprof = rpf.values
+            rad = rpf.rad
+            meta = rpf.meta
         elif name in phyvars.RPROF_EXTRA:
             self._cached_extra[name] = phyvars.RPROF_EXTRA[name](step)
-            rprof, rad, meta = self._cached_extra[name]
+            rpf = self._cached_extra[name]
+            rprof = rpf.values
+            rad = rpf.rad
+            meta = rpf.meta
         else:
             raise error.UnknownRprofVarError(name)
         rprof, _ = step.sdat.scale(rprof, meta.dim)

stagpy/datatypes.py

@@ -2,13 +2,15 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, NamedTuple
+import typing
+from dataclasses import dataclass
 
-if TYPE_CHECKING:
-    from numpy import ndarray
+if typing.TYPE_CHECKING:
+    from numpy.typing import NDArray
 
 
-class Varf(NamedTuple):
+@dataclass(frozen=True)
+class Varf:
     """Metadata of scalar field.
 
     Attributes:
@@ -21,19 +23,21 @@ class Varf(NamedTuple):
     dim: str
 
 
-class Field(NamedTuple):
+@dataclass(frozen=True)
+class Field:
     """Scalar field and associated metadata.
 
     Attributes:
         values: the field itself.
         meta: the metadata of the field.
     """
 
-    values: ndarray
+    values: NDArray
     meta: Varf
 
 
-class Varr(NamedTuple):
+@dataclass(frozen=True)
+class Varr:
     """Metadata of radial profiles.
 
     Attributes:
@@ -46,12 +50,12 @@ class Varr(NamedTuple):
     """
 
     description: str
-    # Callable[[Step], Tuple[ndarray, ndarray]]]
     kind: str
     dim: str
 
 
-class Rprof(NamedTuple):
+@dataclass(frozen=True)
+class Rprof:
     """Radial profile with associated radius and metadata.
 
     Attributes:
@@ -60,12 +64,13 @@ class Rprof(NamedTuple):
         meta: the metadata of the profile.
     """
 
-    values: ndarray
-    rad: ndarray
+    values: NDArray
+    rad: NDArray
     meta: Varr
 
 
-class Vart(NamedTuple):
+@dataclass(frozen=True)
+class Vart:
     """Metadata of time series.
 
     Attributes:
@@ -82,7 +87,8 @@ class Vart(NamedTuple):
     dim: str
 
 
-class Tseries(NamedTuple):
+@dataclass(frozen=True)
+class Tseries:
     """A time series with associated time and metadata.
 
     Attributes:
@@ -91,6 +97,6 @@ class Tseries(NamedTuple):
         meta: the metadata of the series.
     """
 
-    values: ndarray
-    time: ndarray
+    values: NDArray
+    time: NDArray
     meta: Vart

stagpy/field.py

@@ -34,7 +34,7 @@
 
 def _threed_extract(
     step: Step, var: str, walls: bool = False
-) -> Tuple[Tuple[ndarray, ndarray], Any]:
+) -> Tuple[Tuple[ndarray, ndarray], ndarray]:
     """Return suitable slices and coords for 3D fields."""
     is_vector = not valid_field_var(var)
     hwalls = is_vector or walls
@@ -101,24 +101,28 @@ def get_meshes_fld(
         x position, y position, the values and the metadata of the requested
         field.
     """
-    fld, meta = step.fields[var]
-    hwalls = walls or fld.shape[0] != step.geom.nxtot or fld.shape[1] != step.geom.nytot
+    fld = step.fields[var]
+    hwalls = (
+        walls
+        or fld.values.shape[0] != step.geom.nxtot
+        or fld.values.shape[1] != step.geom.nytot
+    )
     if step.geom.threed and step.geom.cartesian:
-        (xcoord, ycoord), fld = _threed_extract(step, var, walls)
+        (xcoord, ycoord), vals = _threed_extract(step, var, walls)
     elif step.geom.twod_xz:
         xcoord = step.geom.x_walls if hwalls else step.geom.x_centers
         ycoord = step.geom.z_walls if walls else step.geom.z_centers
-        fld = fld[:, 0, :, 0]
+        vals = fld.values[:, 0, :, 0]
     else:  # twod_yz
         xcoord = step.geom.y_walls if hwalls else step.geom.y_centers
         ycoord = step.geom.z_walls if walls else step.geom.z_centers
         if step.geom.curvilinear:
             pmesh, rmesh = np.meshgrid(xcoord, ycoord, indexing="ij")
             xmesh, ymesh = rmesh * np.cos(pmesh), rmesh * np.sin(pmesh)
-        fld = fld[0, :, :, 0]
+        vals = fld.values[0, :, :, 0]
     if step.geom.cartesian:
         xmesh, ymesh = np.meshgrid(xcoord, ycoord, indexing="ij")
-    return xmesh, ymesh, fld, meta
+    return xmesh, ymesh, vals, fld.meta
 
 
 def get_meshes_vec(step: Step, var: str) -> Tuple[ndarray, ndarray, ndarray, ndarray]:
@@ -336,15 +340,15 @@ def _findminmax(
     for step in sdat.walk.filter(snap=True):
         for var in sovs:
             if var in step.fields:
-                field, meta = step.fields[var]
-                field, _ = sdat.scale(field, meta.dim)
+                field = step.fields[var]
+                vals, _ = sdat.scale(field.values, field.meta.dim)
                 if var in minmax:
                     minmax[var] = (
-                        min(minmax[var][0], np.nanmin(field)),
-                        max(minmax[var][1], np.nanmax(field)),
+                        min(minmax[var][0], np.nanmin(vals)),
+                        max(minmax[var][1], np.nanmax(vals)),
                     )
                 else:
-                    minmax[var] = np.nanmin(field), np.nanmax(field)
+                    minmax[var] = np.nanmin(vals), np.nanmax(vals)
     return minmax
 
 

stagpy/plates.py

@@ -166,8 +166,8 @@ def _surf_diag(snap: Step, name: str) -> Field:
         return snap.sfields[name]
     isurf = _isurf(snap)
     with suppress(error.UnknownVarError):
-        field, meta = snap.fields[name]
-        return Field(field[0, :, isurf, 0], meta)
+        field = snap.fields[name]
+        return Field(field.values[0, :, isurf, 0], field.meta)
     if name == "dv2":
         vphi = snap.fields["v2"].values[0, :, isurf, 0]
         if snap.geom.cartesian:
@@ -229,14 +229,14 @@ def plot_at_surface(snap: Step, names: Sequence[Sequence[Sequence[str]]]) -> Non
             fname += "_".join(vplt) + "_"
             label = ""
             for name in vplt:
-                data, meta = _surf_diag(snap, name)
-                label = meta.description
+                field = _surf_diag(snap, name)
+                label = field.meta.description
                 phi = (
                     snap.geom.p_centers
-                    if data.size == snap.geom.nptot
+                    if field.values.size == snap.geom.nptot
                     else snap.geom.p_walls
                 )
-                axis.plot(phi, data, label=label)
+                axis.plot(phi, field.values, label=label)
                 axis.set_ylim([conf.plot.vmin, conf.plot.vmax])
             if conf.plates.continents:
                 continents = _continents_location(snap)

stagpy/processing.py

@@ -50,7 +50,9 @@ def dt_dt(sdat: StagyyData) -> Tseries:
     Returns:
         derivative of temperature and time arrays.
     """
-    temp, time, _ = sdat.tseries["Tmean"]
+    series = sdat.tseries["Tmean"]
+    temp = series.values
+    time = series.time
     dtdt = (temp[1:] - temp[:-1]) / (time[1:] - time[:-1])
     return Tseries(dtdt, time[:-1], Vart("Derivative of temperature", r"dT/dt", "K/s"))
 
@@ -73,12 +75,12 @@ def ebalance(sdat: StagyyData) -> Tseries:
     else:
         coefsurf = 1.0
         volume = 1.0
-    dtdt, time, _ = dt_dt(sdat)
+    dtdt = 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 Tseries(ebal, time, Vart("Energy balance", r"$\mathrm{Nu}$", "1"))
+    ebal = ftop[1:] - fbot[1:] + volume * (dtdt.values - radio[1:])
+    return Tseries(ebal, dtdt.time, Vart("Energy balance", r"$\mathrm{Nu}$", "1"))
 
 
 def mobility(sdat: StagyyData) -> Tseries:
@@ -154,9 +156,9 @@ def diffs_prof(step: Step) -> Rprof:
     Returns:
         the diffusion and radius.
     """
-    diff, rad, _ = diff_prof(step)
+    rpf = diff_prof(step)
     meta = Varr("Scaled diffusion", "Heat flux", "W/m2")
-    return Rprof(_scale_prof(step, diff, rad), rad, meta)
+    return Rprof(_scale_prof(step, rpf.values, rpf.rad), rpf.rad, meta)
 
 
 def advts_prof(step: Step) -> Rprof:
@@ -220,10 +222,12 @@ def energy_prof(step: Step) -> Rprof:
     Returns:
         the energy flux and radius.
     """
-    diff, rad, _ = diffs_prof(step)
-    adv, _, _ = advts_prof(step)
+    diff_p = diffs_prof(step)
+    adv_p = advts_prof(step)
     return Rprof(
-        diff + np.append(adv, 0), rad, Varr("Total heat flux", "Heat flux", "W/m2")
+        diff_p.values + np.append(adv_p.values, 0),
+        diff_p.rad,
+        Varr("Total heat flux", "Heat flux", "W/m2"),
     )
 
 

stagpy/rprof.py

@@ -36,7 +36,10 @@ def plot_rprofs(rprofs: _Rprofs, names: Sequence[Sequence[Sequence[str]]]) -> No
             xlabel = None
             profs_on_plt = (rprofs[rvar] for rvar in vplt)
             fname += "_".join(vplt) + "_"
-            for ivar, (rprof, rad, meta) in enumerate(profs_on_plt):
+            for ivar, rpf in enumerate(profs_on_plt):
+                rprof = rpf.values
+                rad = rpf.rad
+                meta = rpf.meta
                 if conf.rprof.depth:
                     rad = sdat.scale(rprofs.bounds[1], "m")[0] - rad
                 axes[iplt].plot(rprof, rad, conf.rprof.style, label=meta.description)
@@ -73,16 +76,16 @@ def plot_grid(step: Step) -> None:
         step (:class:`~stagpy._step.Step`): a step of a StagyyData
             instance.
     """
-    drad, rad, _ = step.rprofs["dr"]
+    drprof = step.rprofs["dr"]
     _, unit = step.sdat.scale(1, "m")
     if unit:
         unit = f" ({unit})"
     fig, (ax1, ax2) = plt.subplots(2, sharex=True)
-    ax1.plot(rad, "-ko")
+    ax1.plot(drprof.rad, "-ko")
     ax1.set_ylabel("$r$" + unit)
-    ax2.plot(drad, "-ko")
+    ax2.plot(drprof.values, "-ko")
     ax2.set_ylabel("$dr$" + unit)
-    ax2.set_xlim([-0.5, len(rad) - 0.5])
+    ax2.set_xlim([-0.5, len(drprof.rad) - 0.5])
     ax2.set_xlabel("Cell number")
     _helpers.saveplot(fig, "grid", step.istep)
 

stagpy/stagyydata.py

@@ -267,10 +267,16 @@ def __getitem__(self, name: str) -> Tseries:
             else:
                 meta = Vart(name, "", "1")
         elif name in self._cached_extra:
-            series, time, meta = self._cached_extra[name]
+            tseries = self._cached_extra[name]
+            series = tseries.values
+            time = tseries.time
+            meta = tseries.meta
         elif name in phyvars.TIME_EXTRA:
             self._cached_extra[name] = phyvars.TIME_EXTRA[name](self.sdat)
-            series, time, meta = self._cached_extra[name]
+            tseries = self._cached_extra[name]
+            series = tseries.values
+            time = tseries.time
+            meta = tseries.meta
         else:
             raise error.UnknownTimeVarError(name)
         series, _ = self.sdat.scale(series, meta.dim)
@@ -289,14 +295,18 @@ def tslice(
             tend: ending time. Set to None to stop at the end of available
                 data.
         """
-        data, time, meta = self[name]
+        series = self[name]
         istart = 0
-        iend = len(time)
+        iend = len(series.time)
         if tstart is not None:
-            istart = _helpers.find_in_sorted_arr(tstart, time)
+            istart = _helpers.find_in_sorted_arr(tstart, series.time)
         if tend is not None:
-            iend = _helpers.find_in_sorted_arr(tend, time, True) + 1
-        return Tseries(data[istart:iend], time[istart:iend], meta)
+            iend = _helpers.find_in_sorted_arr(tend, series.time, True) + 1
+        return Tseries(
+            series.values[istart:iend],
+            series.time[istart:iend],
+            series.meta,
+        )
 
     @property
     def time(self) -> ndarray:
@@ -342,14 +352,14 @@ def __getitem__(self, name: str) -> Rprof:
         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)
+        rpf = next(steps_iter).rprofs[name]
+        rprof = np.copy(rpf.values)
         nprofs = 1
         for step in steps_iter:
             nprofs += 1
             rprof += step.rprofs[name].values
         rprof /= nprofs
-        self._cached_data[name] = Rprof(rprof, rad, meta)
+        self._cached_data[name] = Rprof(rprof, rpf.rad, rpf.meta)
         return self._cached_data[name]
 
     @property

stagpy/stagyyparsers.py

@@ -9,6 +9,7 @@
 
 import re
 import typing
+from dataclasses import dataclass
 from functools import partial
 from itertools import product
 from operator import itemgetter
@@ -368,14 +369,15 @@ def _readbin(
     return elts
 
 
-class _HeaderInfo(typing.NamedTuple):
+@dataclass(frozen=True)
+class _HeaderInfo:
     """Header information."""
 
     magic: int
     nval: int
     sfield: bool
     readbin: Callable
-    header: Dict[str, Any]
+    header: dict[str, Any]
 
 
 def _legacy_header(