mv StagyyData.walk to _helpers.walk

Author: amorison

stagpy/_helpers.py

@@ -15,6 +15,21 @@
     from matplotlib.figure import Figure
     from numpy import ndarray
 
+    from .config import Config
+    from .stagyydata import StagyyData, _StepsView
+
+
+def walk(sdat: StagyyData, conf: Config) -> _StepsView:
+    """Return view on configured steps slice.
+
+    Other Parameters:
+        conf.core.snapshots: the slice of snapshots.
+        conf.core.timesteps: the slice of timesteps.
+    """
+    if conf.core.timesteps:
+        return sdat.steps[conf.core.timesteps]
+    return sdat.snaps[conf.core.snapshots]
+
 
 def out_name(stem: str, timestep: Optional[int] = None) -> str:
     """Return StagPy out file name.

stagpy/commands.py

@@ -10,7 +10,7 @@
 from textwrap import TextWrapper, indent
 
 from . import __version__, phyvars, stagyydata
-from ._helpers import baredoc
+from ._helpers import baredoc, walk
 from .config import CONFIG_LOCAL, Config
 
 if typing.TYPE_CHECKING:
@@ -46,7 +46,7 @@ def info_cmd() -> None:
     else:
         print("Spherical", dimension)
     print()
-    for step in sdat.walk:
+    for step in walk(sdat, conf):
         print(f"Step {step.istep}/{lstep.istep}", end="")
         if step.isnap is not None:
             print(f", snapshot {step.isnap}/{lsnap.isnap}")

stagpy/field.py

@@ -27,6 +27,7 @@
 
     from ._step import Step
     from .datatypes import Varf
+    from .stagyydata import _StepsView
 
 
 # The location is off for vertical velocities: they have an extra
@@ -335,11 +336,11 @@ def plot_vec(
 
 
 def _findminmax(
-    sdat: StagyyData, sovs: Iterable[str]
+    view: _StepsView, sovs: Iterable[str]
 ) -> Dict[str, Tuple[float, float]]:
     """Find min and max values of several fields."""
     minmax: Dict[str, Tuple[float, float]] = {}
-    for step in sdat.walk.filter(snap=True):
+    for step in view.filter(snap=True):
         for var in sovs:
             if var in step.fields:
                 vals = step.fields[var].values
@@ -363,15 +364,16 @@ def cmd() -> None:
     from . import conf
 
     sdat = StagyyData(conf.core.path)
+    view = _helpers.walk(sdat, conf)
     # no more than two fields in a subplot
     lovs = [[slov[:2] for slov in plov] for plov in conf.field.plot]
     minmax = {}
     if conf.plot.cminmax:
         conf.plot.vmin = None
         conf.plot.vmax = None
         sovs = set(slov[0] for plov in lovs for slov in plov)
-        minmax = _findminmax(sdat, sovs)
-    for step in sdat.walk.filter(snap=True):
+        minmax = _findminmax(view, sovs)
+    for step in view.filter(snap=True):
         for vfig in lovs:
             fig, axes = plt.subplots(
                 ncols=len(vfig), squeeze=False, figsize=(6 * len(vfig), 6)

stagpy/plates.py

@@ -400,21 +400,22 @@ def cmd() -> None:
     from . import conf
 
     sdat = StagyyData(conf.core.path)
+    view = _helpers.walk(sdat, conf)
 
-    isurf = _isurf(next(iter(sdat.walk)))
-    vrms_surf = sdat.walk.filter(rprofs=True).rprofs_averaged["vhrms"].values[isurf]
+    isurf = _isurf(next(iter(view)))
+    vrms_surf = view.filter(rprofs=True).rprofs_averaged["vhrms"].values[isurf]
     nb_plates = []
     time = []
     istart, iend = None, None
 
-    oname = _helpers.out_name(f"plates_trenches_{sdat.walk.stepstr}")
+    oname = _helpers.out_name(f"plates_trenches_{view.stepstr}")
     with open(f"{oname}.dat", "w") as fid:
         fid.write(
             "#  istep     time   time_My   phi_trench  vel_trench  "
             "distance     phi_cont  age_trench_My\n"
         )
 
-        for step in sdat.walk.filter(fields=["T"]):
+        for step in view.filter(fields=["T"]):
             # could check other fields too
             _write_trench_diagnostics(step, vrms_surf, fid, conf)
             plot_at_surface(step, conf.plates.plot, conf)

stagpy/rprof.py

@@ -99,13 +99,14 @@ def cmd() -> None:
     from . import conf
 
     sdat = StagyyData(conf.core.path)
+    view = _helpers.walk(sdat, conf)
 
     if conf.rprof.grid:
-        for step in sdat.walk.filter(rprofs=True):
+        for step in view.filter(rprofs=True):
             plot_grid(step)
 
     if conf.rprof.average:
-        plot_rprofs(sdat.walk.rprofs_averaged, conf.rprof.plot, conf)
+        plot_rprofs(view.rprofs_averaged, conf.rprof.plot, conf)
     else:
-        for step in sdat.walk.filter(rprofs=True):
+        for step in view.filter(rprofs=True):
             plot_rprofs(step.rprofs, conf.rprof.plot, conf)

stagpy/stagyydata.py

@@ -19,7 +19,7 @@
 
 import numpy as np
 
-from . import _helpers, _step, conf, error, phyvars, stagyyparsers
+from . import _helpers, _step, error, phyvars, stagyyparsers
 from ._step import Step
 from .datatypes import Rprof, Tseries, Vart
 from .parfile import StagyyPar
@@ -762,18 +762,6 @@ def _files(self) -> Set[Path]:
             return set(out_dir.iterdir())
         return set()
 
-    @property
-    def walk(self) -> _StepsView:
-        """Return view on configured steps slice.
-
-        Other Parameters:
-            conf.core.snapshots: the slice of snapshots.
-            conf.core.timesteps: the slice of timesteps.
-        """
-        if conf.core.timesteps:
-            return self.steps[conf.core.timesteps]
-        return self.snaps[conf.core.snapshots]
-
     @property
     def nfields_max(self) -> Optional[int]:
         """Maximum number of scalar fields kept in memory.

tests/test_helpers.py

@@ -1,5 +1,18 @@
+import pytest
+
 import stagpy
 import stagpy._helpers
+from stagpy import _helpers
+from stagpy.config import Config
+from stagpy.stagyydata import StagyyData
+
+
+def test_walk_dflt(sdat: StagyyData) -> None:
+    view = _helpers.walk(sdat, Config.default_())
+    wlk = iter(view)
+    assert next(wlk) is sdat.snaps[-1]
+    with pytest.raises(StopIteration):
+        next(wlk)
 
 
 def test_out_name_conf() -> None:

tests/test_stagyydata.py

@@ -35,13 +35,6 @@ def test_sdat_tseries(sdat: StagyyData) -> None:
     assert sdat.tseries["Tmean"].time is sdat.tseries.time
 
 
-def test_sdat_walk_dflt(sdat: StagyyData) -> None:
-    wlk = iter(sdat.walk)
-    assert next(wlk) is sdat.snaps[-1]
-    with pytest.raises(StopIteration):
-        next(wlk)
-
-
 def test_steps_iter(sdat: StagyyData) -> None:
     assert sdat.steps[:] == sdat.steps