Add generic dump function for IO_Params.

hysop/tools/io_utils.py

@@ -49,8 +49,8 @@ class IO(object):
     @requires_cmd('stat')
     def get_fs_type(path):
         cmd = ['stat', '-f', '-c', '%T', path]
-        fs_type=''
-        if mpi.main_rank==0:
+        fs_type = ''
+        if mpi.main_rank == 0:
             fs_type = subprocess.check_output(cmd)
         fs_type = mpi.main_comm.bcast(fs_type, root=0)
         return fs_type.replace('\n', '')
@@ -151,7 +151,7 @@ class IO(object):
             raise RuntimeError(msg)
         if not os.path.exists(cpath):
             try:
-                if mpi.main_rank==0:
+                if mpi.main_rank == 0:
                     os.makedirs(cpath)
             except IOError:
                 pass
@@ -232,7 +232,7 @@ class IO(object):
 class IOParams(namedtuple("IOParams", ['filename', 'filepath',
                                        'frequency', 'fileformat',
                                        'dump_times_fp32', 'dump_times_fp64',
-                                       'dump_tstart', 'dump_tend',
+                                       'dump_tstart', 'dump_tend', 'dump_func',
                                        'io_leader', 'visu_leader', 'with_last',
                                        'enable_ram_fs', 'force_ram_fs',
                                        'dump_is_temporary', 'postprocess_dump',
@@ -254,8 +254,10 @@ class IOParams(namedtuple("IOParams", ['filename', 'filepath',
         Extra dump times that should be used to dump in addition to frequency (double precision)
     dump_tstart: float
         Start to dump at given time. Defaults to -np.inf (no time constraints).
-    dump_end: float
+    dump_tend: float
         Stop to dump at given time. Defaults to +np.inf (no time constraints).
+    dump_func: function
+        Generic function to compute the should_dump result.
     with_last: boolean
         should dump when iteration is last one
     io_leader : int
@@ -290,7 +292,7 @@ class IOParams(namedtuple("IOParams", ['filename', 'filepath',
     """
     def __new__(cls, filename, filepath=None,
                 frequency=1, fileformat=None,
-                dump_times=None,  dump_tstart=None, dump_tend=None,
+                dump_times=None,  dump_tstart=None, dump_tend=None, dump_func=None,
                 io_leader=0, visu_leader=0, with_last=False,
                 enable_ram_fs=False, force_ram_fs=False,
                 dump_is_temporary=False, postprocess_dump=None,
@@ -316,7 +318,9 @@ class IOParams(namedtuple("IOParams", ['filename', 'filepath',
         check_instance(dump_is_temporary, bool)
         check_instance(postprocess_dump, str, allow_none=True)
         check_instance(hdf5_disable_compression, bool)
-
+        if dump_func:
+            assert callable(dump_func), "given function must be callable"
+            assert dump_func.func_code.co_argcount, "given function must take one arg (as simulation object)"
         frequency = int(frequency)
         dump_tstart = float(dump_tstart)
         dump_tend = float(dump_tend)
@@ -378,13 +382,15 @@ class IOParams(namedtuple("IOParams", ['filename', 'filepath',
         return super(IOParams, cls).__new__(cls, filename, filepath,
                                             frequency, fileformat,
                                             dump_times_fp32, dump_times_fp64,
-                                            dump_tstart, dump_tend,
+                                            dump_tstart, dump_tend, dump_func,
                                             io_leader, visu_leader, with_last,
                                             enable_ram_fs, force_ram_fs,
                                             dump_is_temporary, postprocess_dump,
                                             hdf5_disable_compression, disk_filepath, kwds)
 
     def should_dump(self, simulation):
+        if self.dump_func is not None:
+            return self.dump_func(simulation)
         frequency = self.frequency
         t = simulation.t()
         dump = (frequency >= 0) and (self.with_last and simulation._next_is_last)
@@ -403,7 +409,7 @@ class IOParams(namedtuple("IOParams", ['filename', 'filepath',
 
     def clone(self, **kwds):
         keys = ('filename', 'frequency', 'fileformat',
-                'dump_times', 'dump_tstart', 'dump_tend',
+                'dump_times', 'dump_tstart', 'dump_tend', 'dump_func',
                 'io_leader', 'visu_leader', 'with_last',
                 'enable_ram_fs', 'force_ram_fs',
                 'dump_is_temporary', 'postprocess_dump',
@@ -441,6 +447,7 @@ fileformat:    {}
 frequency:     {}
 dump_times:    {}
 dump_tstart:   {}
+dump_func:     {}
 dump_tend:     {}
 io_leader:     {}
 visu_leader:   {}
@@ -451,7 +458,7 @@ post_process:  {}
 hdf5_no_compr: {}
 extra_kwds:    {}'''.format(
                 self.filename, self.filepath, self.fileformat,
-                self.frequency, self.dump_times, self.dump_tstart, self.dump_tend,
+                self.frequency, self.dump_times, self.dump_tstart, self.dump_tend, self.dump_func,
                 self.io_leader, self.visu_leader,
                 self.enable_ram_fs, self.force_ram_fs,
                 self.dump_is_temporary,