Merge pull request #87 from StagPython/xmf

Improve performances when reading tracers and fields with HDF5 output

Author: amorison

stagpy/_step.py

@@ -367,16 +367,15 @@ def _get_raw_data(self, name: str) -> Tuple[List[str], Any]:
                 (stem, fvars) for stem, fvars in self._filesh5.items() if name in fvars
             ]
             for filestem, list_fvar in files:
+                sdat = self.step.sdat
                 if filestem in phyvars.SFIELD_FILES_H5:
-                    xmff = "Data{}.xmf".format(
-                        "Bottom" if name.endswith("bot") else "Surface"
-                    )
+                    xmff = sdat._botxmf if name.endswith("bot") else sdat._topxmf
                     header = self._header
                 else:
-                    xmff = "Data.xmf"
+                    xmff = sdat._dataxmf
                     header = None
                 parsed_data = stagyyparsers.read_field_h5(
-                    self.step.sdat.hdf5 / xmff, filestem, self.step.isnap, header
+                    xmff, filestem, self.step.isnap, header
                 )
                 if parsed_data is not None:
                     break
@@ -405,8 +404,9 @@ def _header(self) -> Optional[Dict[str, Any]]:
         if binfiles:
             header = stagyyparsers.field_header(binfiles.pop())
         elif self.step.sdat.hdf5:
-            xmf = self.step.sdat.hdf5 / "Data.xmf"
-            header = stagyyparsers.read_geom_h5(xmf, self.step.isnap)[0]
+            header = stagyyparsers.read_geom_h5(
+                self.step.sdat._dataxmf, self.step.isnap
+            )
         return header if header else None
 
     @cached_property
@@ -447,14 +447,10 @@ def __getitem__(self, name: str) -> Optional[List[ndarray]]:
             self.step.sdat.filename("tra", timestep=self.step.isnap, force_legacy=True)
         )
         if data is None and self.step.sdat.hdf5:
-            position = any(axis not in self._data for axis in "xyz")
-            self._data.update(
-                stagyyparsers.read_tracers_h5(
-                    self.step.sdat.hdf5 / "DataTracers.xmf",
-                    name,
-                    self.step.isnap,
-                    position,
-                )
+            self._data[name] = stagyyparsers.read_tracers_h5(
+                self.step.sdat._traxmf,
+                name,
+                self.step.isnap,
             )
         elif data is not None:
             self._data.update(data)

stagpy/stagyydata.py

@@ -22,6 +22,7 @@
 from . import _helpers, _step, conf, error, parfile, phyvars, stagyyparsers
 from ._step import Step
 from .datatypes import Rprof, Tseries, Vart
+from .stagyyparsers import FieldXmf, TracersXmf
 
 if typing.TYPE_CHECKING:
     from os import PathLike
@@ -794,6 +795,34 @@ def hdf5(self) -> Optional[Path]:
         h5_folder = self.path / self.par["ioin"]["hdf5_output_folder"]
         return h5_folder if (h5_folder / "Data.xmf").is_file() else None
 
+    @cached_property
+    def _dataxmf(self) -> FieldXmf:
+        assert self.hdf5 is not None
+        return FieldXmf(
+            path=self.hdf5 / "Data.xmf",
+        )
+
+    @cached_property
+    def _topxmf(self) -> FieldXmf:
+        assert self.hdf5 is not None
+        return FieldXmf(
+            path=self.hdf5 / "DataSurface.xmf",
+        )
+
+    @cached_property
+    def _botxmf(self) -> FieldXmf:
+        assert self.hdf5 is not None
+        return FieldXmf(
+            path=self.hdf5 / "DataBottom.xmf",
+        )
+
+    @cached_property
+    def _traxmf(self) -> TracersXmf:
+        assert self.hdf5 is not None
+        return TracersXmf(
+            path=self.hdf5 / "DataTracers.xmf",
+        )
+
     @property
     def par(self) -> Namelist:
         """Content of par file.