rprof parsers use names in output files

Author: amorison

src/stagpy/stagyydata.py

@@ -732,14 +732,14 @@ def par(self) -> StagyyPar:
     @cached_property
     def _rprof_and_times(self) -> tuple[dict[int, DataFrame], DataFrame | None]:
         rproffile = self.filename("rprof.h5")
-        data = stagyyparsers.rprof_h5(rproffile, list(phyvars.RPROF.keys()))
+        data = stagyyparsers.rprof_h5(rproffile)
         if data[1] is not None:
             return data
         rproffile = self.filename("rprof.dat")
         if self.hdf5 and not rproffile.is_file():
             # check legacy folder as well
             rproffile = self.filename("rprof.dat", force_legacy=True)
-        return stagyyparsers.rprof(rproffile, list(phyvars.RPROF.keys()))
+        return stagyyparsers.rprof(rproffile)
 
     @property
     def rtimes(self) -> DataFrame | None:

src/stagpy/stagyyparsers.py

@@ -20,7 +20,7 @@
 import pandas as pd
 
 from .error import ParsingError
-from .phyvars import FIELD_FILES_H5, SFIELD_FILES_H5
+from .phyvars import FIELD_FILES_H5, RPROF, SFIELD_FILES_H5
 from .xdmf import XmlStream
 
 if typing.TYPE_CHECKING:
@@ -165,25 +165,24 @@ def _extract_rsnap_isteps(
     return isteps
 
 
-def rprof(
-    rproffile: Path, colnames: list[str]
-) -> tuple[dict[int, DataFrame], DataFrame | None]:
+def rprof(rproffile: Path) -> tuple[dict[int, DataFrame], DataFrame | None]:
     """Extract radial profiles data.
 
-    If `colnames` is too long, it will be truncated. If it is too short,
-    additional numeric column names from 0 to N-1 will be attributed to the N
-    extra columns present in `timefile`.
-
     Args:
         rproffile: path of the rprof.dat file.
-        colnames: names of the variables expected in `rproffile`.
 
     Returns:
         profs: a dict mapping istep to radial profiles.
         times: the time indexed by time steps.
     """
     if not rproffile.is_file():
         return {}, None
+
+    with rproffile.open() as fid:
+        colnames = fid.readline().strip().split()
+    if not colnames:
+        colnames = list(RPROF.keys())
+
     data = pd.read_csv(
         rproffile,
         sep=r"\s+",
@@ -213,17 +212,11 @@ def rprof(
     return all_data, df_times
 
 
-def rprof_h5(
-    rproffile: Path, colnames: list[str]
-) -> tuple[dict[int, DataFrame], DataFrame | None]:
+def rprof_h5(rproffile: Path) -> tuple[dict[int, DataFrame], DataFrame | None]:
     """Extract radial profiles data.
 
-    If `colnames` is too long, it will be truncated. If it is too short,
-    additional column names will be deduced from the content of the file.
-
     Args:
         rproffile: path of the rprof.h5 file.
-        colnames: names of the variables expected in `rproffile`.
 
     Returns:
         profs: a dict mapping istep to radial profiles.
@@ -235,13 +228,13 @@ def rprof_h5(
     data = {}
     with h5py.File(rproffile, "r") as h5f:
         dnames = sorted(dname for dname in h5f.keys() if dname.startswith("rprof_"))
-        h5names = h5f["names"].asstr()[len(colnames) :]
+        colnames = h5f["names"].asstr()[()]
         for dname in dnames:
             dset = h5f[dname]
             arr = dset[()]
             istep = dset.attrs["istep"]
             step_cols = list(colnames)
-            _tidy_names(step_cols, arr.shape[1], h5names)  # check shape
+            _tidy_names(step_cols, arr.shape[1])  # check shape
             data[istep] = pd.DataFrame(arr, columns=step_cols)
             isteps.append((istep, dset.attrs["time"]))
 

tests/test_parsers.py

@@ -25,25 +25,20 @@ def test_time_series_invalid_prs() -> None:
 
 def test_rprof_prs(sdat_legacy: StagyyData) -> None:
     sdat = sdat_legacy
-    names = ["aa", "bb", "cc"]
-    data, time = prs.rprof(sdat.filename("rprof.dat"), list(names))
-    assert all((df.columns[:3] == names).all() for df in data.values())
-    assert all(
-        (df.columns[3:] == list(map(str, range(df.shape[1] - 3)))).all()
-        for df in data.values()
-    )
+    data, time = prs.rprof(sdat.filename("rprof.dat"))
+    assert all((df.columns[:3] == ["r", "Tmean", "Tmin"]).all() for df in data.values())
 
 
 def test_rprof_h5(sdat_h5: StagyyData) -> None:
     sdat = sdat_h5
     assert sdat.hdf5 is not None
-    data, _times = prs.rprof_h5(sdat.hdf5 / "rprof.h5", colnames=[])
+    data, _times = prs.rprof_h5(sdat.hdf5 / "rprof.h5")
     assert data is not None
     assert (data[1000].columns[:3] == ["r", "Tmean", "Tmin"]).all()
 
 
 def test_rprof_invalid_prs() -> None:
-    assert prs.rprof(Path("dummy"), []) == ({}, None)
+    assert prs.rprof(Path("dummy")) == ({}, None)
 
 
 def test_fields_prs(sdat_legacy: StagyyData) -> None: