Merge branch 'topic/default/util-phys-fields' into 'branch/default'

Better util.phys_fields

See merge request fluiddyn/fluidsim!443

Author: paugier

fluidsim/base/output/model.xmf

@@ -15,17 +15,17 @@
       </Geometry>
       <Attribute Name="vx" AttributeType="Scalar" Center="Cell">
 	<DataItem Dimensions="16 16 16" NumberType="Float" Format="HDF">
-	  state_phys_t=000.000_it=0.hd5:/state_phys/vx
+	  state_phys_t000.000_it0.hd5:/state_phys/vx
 	</DataItem>
       </Attribute>
       <Attribute Name="vy" AttributeType="Scalar" Center="Cell">
 	<DataItem Dimensions="16 16 16" NumberType="Float" Format="HDF">
-	  state_phys_t=000.000_it=0.hd5:/state_phys/vy
+	  state_phys_t000.000_it0.hd5:/state_phys/vy
 	</DataItem>
       </Attribute>
       <Attribute Name="vz" AttributeType="Scalar" Center="Cell">
 	<DataItem Dimensions="16 16 16" NumberType="Float" Format="HDF">
-	  state_phys_t=000.000_it=0.hd5:/state_phys/vz
+	  state_phys_t000.000_it0.hd5:/state_phys/vz
 	</DataItem>
       </Attribute>
     </Grid>

fluidsim/util/phys_fields.py

@@ -147,27 +147,46 @@ def save_file(
 
 def compute_file_name(time, str_width, ext, it=None):
     """Compute the file name from time and co"""
-    str_it = "" if it is None else f"_{it=}"
+    str_it = "" if it is None else f"_it{it}"
     return f"state_phys_t{time:0{str_width}.3f}{str_it}.{ext}"
 
 
-def time_from_path(path):
+# Module-level variable, compiled on first use
+_TIME_PATTERN = None
+
+
+def time_from_path(path, exact=False):
     """Regular expression search to extract time from filename."""
+
+    if exact:
+        with h5py.File(path, "r") as file:
+            time = file.attrs["time"]
+        return time
+
+    global _TIME_PATTERN
+    if _TIME_PATTERN is None:
+        _TIME_PATTERN = re.compile(
+            r"""
+            (?!t)     # text after t but exclude it
+            [0-9]+    # a couple of digits
+            \.        # the decimal point
+            [0-9]+    # a couple of digits
+            """,
+            re.VERBOSE,
+        )
+
     filename = os.path.basename(path)
-    pattern = r"""
-        (?!t)     # text after t but exclude it
-        [0-9]+    # a couple of digits
-        \.        # the decimal point
-        [0-9]+    # a couple of digits
-    """
-    match = re.search(pattern, filename, re.VERBOSE)
+    match = _TIME_PATTERN.search(filename)
     time = float(match.group(0))
     return time
 
 
 def name_file_from_time_approx(path_dir, t_approx=None):
     """Return the file name whose time is the closest to the given time.
 
+    Warning: for parallel runs and if ``t_approx is not None``, it is safer
+    to only call this function by one process.
+
     Parameters
     ----------
 
@@ -178,11 +197,8 @@ def name_file_from_time_approx(path_dir, t_approx=None):
     t_approx : number or "last" (optional)
 
       Approximate time of the file to be loaded.
-
-    .. todo::
-
-        Can be elegantly implemented using regex as done in
-        ``fluidsim.base.output.phys_fields.time_from_path``
+      If "last", use the last time.
+      If None, just return the last file name (sorted in alphabetic order).
 
     """
     if not isinstance(path_dir, Path):
@@ -198,20 +214,13 @@ def name_file_from_time_approx(path_dir, t_approx=None):
         # should be the last one but not 100% sure
         return path_files[-1].name
 
-    name_files = [path.name for path in path_files]
-    if "state_phys_t=" in name_files[0]:
-        ind_start_time = len("state_phys_t=")
-    else:
-        ind_start_time = len("state_phys_t")
-
-    times = np.empty([nb_files])
-    for ii, name in enumerate(name_files):
-        tmp = ".".join(name[ind_start_time:].split(".")[:2])
-        if "_" in tmp:
-            tmp = tmp[: tmp.index("_")]
-        times[ii] = float(tmp)
+    # the time are read from the files if at least one of the name contains "_it"
+    exact = any("_it" in path.name for path in path_files)
+    times = [time_from_path(path, exact=exact) for path in path_files]
+
     if t_approx == "last":
-        t_approx = times.max()
-    i_file = abs(times - t_approx).argmin()
-    name_file = path_files[i_file].name
-    return name_file
+        path_file = max(zip(times, path_files))[1]
+    else:
+        i_file = abs(np.array(times) - t_approx).argmin()
+        path_file = path_files[i_file]
+    return path_file.name

fluidsim/util/test_phys_fields.py

@@ -1,19 +1,25 @@
+import pytest
+
 import numpy as np
 
 import h5py
 
 from fluidsim.util.phys_fields import (
     name_file_from_time_approx,
     compute_file_name,
+    time_from_path,
 )
 
+from fluiddyn.util import mpi
+
+
+dt = 0.0004
+times = dt * np.arange(11)
 
-def test_name_file_from_time_approx(tmp_path):
-    dt = 0.0004
-    times = dt * np.arange(11)
 
-    path_dir = tmp_path / "run_dir"
-    path_dir.mkdir()
+@pytest.fixture(scope="module")
+def path_dir_with_files(tmp_path_factory):
+    path_dir = tmp_path_factory.mktemp("run_dir")
 
     str_width, ext = 7, "h5"
 
@@ -27,9 +33,29 @@ def test_name_file_from_time_approx(tmp_path):
             file.attrs["time"] = time
             file.attrs["it"] = it
 
-        # print(path_file)
+    return path_dir
+
+
+def test_time_from_path(path_dir_with_files):
+    paths = sorted(path_dir_with_files.glob("*.h5"))
+    for it, path in enumerate(paths):
+        t_from_path = time_from_path(path)
+        assert t_from_path == round(it * dt, 3)
+        t_exact_from_path = time_from_path(path, exact=True)
+        assert t_exact_from_path == it * dt
+
+
+def test_name_file_from_time_approx(path_dir_with_files):
+    if mpi.rank > 0:
+        return
+
+    path_dir = path_dir_with_files
 
     name_file_last = name_file_from_time_approx(path_dir)
-    assert name_file_last == path_file.name
+    assert name_file_last == "state_phys_t000.004_it10.h5"
+
+    name_file = name_file_from_time_approx(path_dir, t_approx=0.0034)
+    assert name_file == "state_phys_t000.003_it8.h5", name_file
 
-    # TODO: add assert statements to get bugs!
+    name_file_last = name_file_from_time_approx(path_dir, t_approx="last")
+    assert name_file_last == "state_phys_t000.004_it10.h5"

fluidsim/util/util.py

@@ -279,10 +279,16 @@ def _path_file_from_time_approx(thing, t_approx):
     if thing is not None and Path(thing).is_file():
         path_file = Path(thing)
     else:
-        path_dir = pathdir_from_namedir(thing)
+        path_dir = Path(pathdir_from_namedir(thing))
         # choose the file with the time closer to t_approx
-        name_file = name_file_from_time_approx(path_dir, t_approx)
-        path_file = Path(path_dir) / name_file
+        if mpi.rank == 0:
+            name_file = name_file_from_time_approx(path_dir, t_approx)
+        else:
+            name_file = None
+        if mpi.nb_proc > 1:
+            name_file = mpi.comm.bcast(name_file, root=0)
+
+        path_file = path_dir / name_file
     return path_file
 
 
@@ -593,7 +599,7 @@ def times_start_last_from_path(path):
         if last_line.startswith("save state_phys"):
             name_file = last_line.split()[-1]
             name_file, ext = os.path.splitext(name_file)
-            word = name_file.split("_it=")[0].split("state_phys_t")[-1]
+            word = name_file.split("_it")[0].split("state_phys_t")[-1]
             t_last = float(word.replace(ext, ""))
         else:
             words = line_it.split()