better io energy handling

examples/taylor_green/taylor_green.py

@@ -136,10 +136,10 @@ def compute(args):
                             variables={velo:npts, vorti: npts}, 
                             projection=args.reprojection_frequency,
                             diffusion=viscosity, dt=dt,
-                            plot_velocity_energy=IOParams(filepath=spectral_path, 
-                                filename='E_velocity_{ite}', frequency=args.dump_freq),
-                            plot_vorticity_energy=IOParams(filepath=spectral_path, 
-                                filename='E_vorticity_{ite}', frequency=args.dump_freq),
+                            dump_energy=IOParams(filepath=spectral_path, filename='E_{fname}.txt', frequency=args.dump_freq),
+                            plot_energy=IOParams(filepath=spectral_path, filename='E_{fname}_{ite}', frequency=args.dump_freq),
+                            plot_input_vorticity_energy=-1,  # <= disable a specific plot
+                            plot_output_vorticity_energy=-1, # <= disable a specific plot
                             implementation=impl, **extra_op_kwds)
     #> We ask to dump the outputs of this operator
     poisson.dump_outputs(fields=(vorti,), frequency=args.dump_freq, **extra_op_kwds)

hysop/backend/device/opencl/operator/poisson_curl.py

@@ -195,9 +195,5 @@ class OpenClPoissonCurl(SpectralPoissonCurlOperatorBase, OpenClSymbolic):
             evt = Fc()
             evt = Bt(simulation=simulation)
         evt = self.exchange_U_ghosts()
-        
-        if self.do_compute_velocity_energy:
-            self.compute_velocity_energy_and_plot(simulation=simulation)
-        if self.do_compute_vorticity_energy:
-            self.compute_vorticity_energy_and_plot(simulation=simulation)
+        self.update_energy(simulation=simulation)        
 

hysop/backend/host/python/operator/poisson_curl.py

@@ -220,9 +220,4 @@ class PythonPoissonCurl(SpectralPoissonCurlOperatorBase, OpenClMappable, HostOpe
         self.dU.exchange_ghosts()
         if (diffuse or project):
             self.dW.exchange_ghosts()
-
-        if self.do_compute_velocity_energy:
-            self.compute_velocity_energy_and_plot(simulation=simulation)
-        if self.do_compute_vorticity_energy:
-            self.compute_vorticity_energy_and_plot(simulation=simulation)
-
+        self.update_energy(simulation=simulation) 

hysop/simulation.py

@@ -373,13 +373,15 @@ class Simulation(object):
     def write_parameters(self, *params, **kwds):
         if ('io_params' not in kwds):
             assert ('filename' in kwds), 'io_params or filename should be specified.'
-            filename   = kwds.pop('filename')
-            filepath   = kwds.pop('filepath',  None)
-            fileformat = kwds.pop('fileformat', IO.ASCII)
-            frequency  = kwds.pop('frequency', 1)
-            io_leader  = kwds.pop('io_leader', 0)
+            filename    = kwds.pop('filename')
+            filepath    = kwds.pop('filepath',  None)
+            fileformat  = kwds.pop('fileformat', IO.ASCII)
+            frequency   = kwds.pop('frequency', 1)
+            io_leader   = kwds.pop('io_leader', 0)
+            visu_leader = kwds.pop('visu_leader', 0)
             io_params = IOParams(filename=filename, filepath=filepath,
-                    frequency=frequency, fileformat=fileformat, io_leader=io_leader)
+                    frequency=frequency, fileformat=fileformat, 
+                    io_leader=io_leader, visu_leader=visu_leader)
         else:
             io_params = kwds.pop('io_params')
         _params = ()
@@ -399,7 +401,7 @@ class Simulation(object):
         return s
     
     def should_dump(self, frequency, with_last=False):
-        dump = (with_last and self._next_is_last)
+        dump = (frequency>=0) and (with_last and self._next_is_last)
         if (frequency >= 0):
             dump |= self.is_time_of_interest
         if (frequency > 0):

hysop/tools/io_utils.py

@@ -176,7 +176,8 @@ class IO(object):
 
 class IOParams(namedtuple("IOParams", ['filename', 'filepath',
                                        'frequency', 'fileformat',
-                                       'io_leader'])):
+                                       'io_leader', 'visu_leader',
+                                       'kwds'])):
     """
     A struct to handle I/O files parameters
 
@@ -192,6 +193,10 @@ class IOParams(namedtuple("IOParams", ['filename', 'filepath',
         format of the file. See notes for available format. Default=HDF5.
     io_leader : int
         rank of the mpi process dealing with the io. Default is 0.
+    visu_leader : int
+        rank of the mpi process dealing with the graphical io. Default is 0.
+    kwds: dict
+        Custom extra keyword arguments to pass to operators
 
     See examples in hysop.operator.hdf_io
 
@@ -203,7 +208,8 @@ class IOParams(namedtuple("IOParams", ['filename', 'filepath',
 
     """
     def __new__(cls, filename, filepath=None, frequency=1,
-                fileformat=None, io_leader=0):
+                fileformat=None, io_leader=0, visu_leader=0,
+                **kwds):
 
         # Filename is absolute path, filepath arg is ignored.
         if os.path.isabs(filename):
@@ -228,7 +234,9 @@ class IOParams(namedtuple("IOParams", ['filename', 'filepath',
 
         IO.check_dir(filename)
         return super(IOParams, cls).__new__(cls, filename, filepath,
-                                            frequency, fileformat, io_leader)
+                                            frequency, fileformat, 
+                                            io_leader, visu_leader,
+                                            kwds)
 
 
 class Writer(object):

hysop/tools/spectral_utils.py

-
+import math, os
 import numpy as np
 import sympy as sm
+import matplotlib
+import matplotlib.pyplot as plt
 
+from hysop import main_rank
+from hysop.tools.io_utils import IOParams
 from hysop.tools.types import check_instance, first_not_None, to_tuple
 from hysop.tools.numerics import is_fp, is_complex, complex_to_float_dtype, float_to_complex_dtype
 from hysop.tools.sympy_utils import Expr, Symbol, Dummy, subscript, tensor_symbol
@@ -704,6 +708,253 @@ def make_trigonometric_polynomial(Xl, Xr, lboundary, rboundary, N):
     return (P, x)
 
 
+class EnergyDumper(object):
+    def __init__(self, energy_parameter, io_params, fname, **kwds):
+        from hysop.parameters.buffer_parameter import BufferParameter
+
+        super(EnergyDumper, self).__init__(**kwds)
+        check_instance(io_params, IOParams)
+        check_instance(io_params.filepath, str)
+        check_instance(io_params.io_leader, int)
+        filename = io_params.filename
+        filename = filename.replace('{fname}', fname)
+        filename = filename.replace('{it}', '')
+        assert '{fname}' not in filename
+        assert '{it}'    not in filename
+        assert io_params.frequency>=0
+
+        check_instance(energy_parameter, BufferParameter)
+        assert energy_parameter.size >= 2
+        assert energy_parameter.dtype in (np.float32, np.float64)
+
+        should_write = (io_params.io_leader == main_rank)
+        
+        if should_write:
+            ulp = np.finfo(energy_parameter.dtype).eps ** 4
+            formatter = {'float_kind':  lambda x: '{:8.8f}'.format(x)}
+
+            if os.path.isfile(filename): 
+                os.remove(filename)
+            f = open(filename, 'a')
+
+            header = u'# Evolution of the power spectrum of {}'.format(energy_parameter.pretty_name.decode('utf-8'))
+            f.write(header.encode('utf-8'))
+            f.write('\n# with mean removed (first coefficient) and values clamped to ulp = epsilon^4 = {}'.format(ulp))
+            f.write('\n# ITERATION  TIME  log10(max(POWER_SPECTRUM[1:], ulp)))')
+            f.flush()
+        else:
+            f = None
+            formatter = None
+            ulp = None
+
+        self.should_write = should_write
+        self.energy_parameter = energy_parameter
+        self.io_params = io_params
+        self.file = f
+        self.formatter = formatter
+        self.ulp = ulp
+       
+    def update(self, simulation, wait_for):
+        if not self.should_write:
+            return
+        if not simulation.should_dump(frequency=self.io_params.frequency, with_last=True):
+            return
+        if (wait_for is not None):
+            wait_for.wait()
+        assert (self.file is not None)
+        energy = self.energy_parameter.value[1:].astype(dtype=np.float64)
+        energy = np.log10(np.maximum(energy, self.ulp))
+        values = np.array2string(energy, 
+                formatter=self.formatter, max_line_width=np.inf)
+        values = '\n{} {} {}'.format(simulation.current_iteration, simulation.t(), values[1:-1])
+        self.file.write(values)
+        self.file.flush()
+
+    @classmethod
+    def do_compute_energy(cls, *args):
+        frequencies = set()
+        for iop in args:
+            if not isinstance(iop, IOParams):
+                continue
+            f = iop.frequency
+            if (f>=0):
+                frequencies.add(f)
+        do_compute = (len(frequencies)>0)
+        frequencies = frequencies if do_compute else None
+        return do_compute, frequencies
+
+    @classmethod
+    def build_energy_parameter(cls, do_compute, field, output_params, prefix):
+        from hysop.parameters.buffer_parameter import BufferParameter
+        if do_compute:
+            pname  = 'E{}_{}'.format(prefix, field.name)
+            pename = 'E{}_{}'.format(prefix, field.pretty_name)
+            param = BufferParameter(name=pname, pretty_name=pename,
+                    shape=None, dtype=None, initial_value=None)
+            assert param.name not in output_params, param.name
+            output_params[param.name] = param
+        else:
+            param = None
+        return param
+
+class EnergyPlotter(object):
+    def __init__(self, energy_parameters, io_params, fname,
+            fig_title=None, axes_shape=(1,), figsize=(15,9),
+            basex=10, basey=10, **kwds):
+        from hysop.parameters.buffer_parameter import BufferParameter
+        super(EnergyPlotter, self).__init__(**kwds)
+        check_instance(io_params, IOParams)
+        check_instance(axes_shape, tuple, minsize=1, allow_none=True)
+        check_instance(energy_parameters, dict, values=BufferParameter)
+
+        should_draw = (io_params.visu_leader == main_rank)
+        filename = io_params.filename
+        filename = filename.replace('{fname}', fname)
+        assert '{ite}' in filename, filename
+        assert io_params.frequency>=0
+        
+        if should_draw:
+            fig, axes = plt.subplots(*axes_shape, figsize=figsize)
+            fig.canvas.mpl_connect('key_press_event', self._on_key_press)
+            fig.canvas.mpl_connect('close_event', self._on_close)
+            
+            axes = np.asarray(axes).reshape(axes_shape)
+            ax = axes[0]
+            
+            if len(energy_parameters)==1:
+                default_fig_title = u'Energy parameter {}'.format(
+                    energy_parameters.values()[0].pretty_name.decode('utf-8'))
+            else:
+                default_fig_title = u'Energy parameters {}'.format(
+                        u' | '.join(p.pretty_name.decode('utf-8') for p in energy_parameters.values()))
+            fig_title = first_not_None(fig_title, default_fig_title)
+            self.fig_title = fig_title
+
+            xmax = 1
+            lines = ()
+            for (label,p) in energy_parameters.iteritems():
+                assert p.size>1
+                Ix = np.arange(1, p.size)
+                xmax = max(xmax, p.size-1)
+                line  = ax.plot(Ix, np.zeros(p.size-1), '--o', label=label, markersize=3)[0]
+                lines += (line,)
+            
+            xmin = 1
+            pmax = math.ceil(math.log(xmax, basex))
+            xmax = basex**pmax if basex==2 else 1.1*xmax
+            ax.set_xlim(xmin, xmax)
+
+            ax.set_title('t=None')
+            ax.set_xlabel('Wavenumber')
+            ax.set_ylabel('Energy')
+            ax.set_xscale('log', basex=basex)
+            ax.set_yscale('log', basey=basey)
+            ax.legend(loc='upper right')
+            ax.grid(True, which='major', ls='--', c='k')
+            ax.grid(True, which='minor', ls=':')
+
+            if (basex == 2):
+                ax.xaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(lambda x, pos: str(int(round(x)))))
+            else:
+                ax.xaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(lambda x, pos: r'$\mathbf{{{}^{{{}}}}}$'.format(basex, int(round(math.log(x, basex))))))
+            ax.yaxis.set_minor_formatter(matplotlib.ticker.FuncFormatter(lambda x, pos, ax=ax: r'$^{{{}}}$'.format(int(round(math.log(x, basey))))))
+            ax.yaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(lambda x, pos: r'$\mathbf{{{}^{{{}}}}}$'.format(basey, int(round(math.log(x, basey))))))
+            
+            running = True
+        else:
+            fig, axes = None, None
+            lines = None
+            running = False
+        
+        ulp = np.finfo(np.find_common_type([], list(p.dtype 
+            for p in energy_parameters.values()))).eps
+        ulp = 1e-4*(ulp**2)
+
+        self.fig  = fig
+        self.axes = axes
+        self.lines = lines
+        
+        self.filename  = filename
+        self.io_params = io_params
+        
+        # see https://stackoverflow.com/questions/8257385/automatic-detection-of-display-availability-with-matplotlib
+        self.has_gui_running = hasattr(fig, 'show')
+
+        self.should_draw = should_draw
+        self.basex = basex
+        self.basey = basey
+
+        self.energy_parameters = energy_parameters
+        self.ulp = ulp
+
+    def update(self, simulation, wait_for):
+        if not self.should_draw:
+            return
+        if not simulation.should_dump(frequency=self.io_params.frequency, with_last=True):
+           return
+        if (wait_for is not None):
+            wait_for.wait()
+        ite = simulation.current_iteration
+        self._update_plot(ite, simulation.t())
+        self._draw()
+        self._savefig(ite)
+
+    def _update_plot(self, iteration, time):
+        ymin=np.PINF
+        ymax=np.NINF
+        for (p,line) in zip(self.energy_parameters.values(), self.lines):
+            energy = p.value[1:]
+            energy = np.maximum(energy, self.ulp)
+            ymin = min(ymin, np.min(energy))
+            ymax = max(ymax, np.max(energy))
+            line.set_ydata(energy)
+
+        basey = self.basey
+        pmin = int(math.floor(math.log(ymin, basey)))
+        pmax = int(math.ceil(math.log(ymax, basey)))
+        if (pmax==pmin):
+            pmax += 1
+        max_majors = 8
+        nlevels  = pmax-pmin+1
+        nsubint = 1
+        while(math.ceil(nlevels/float(nsubint))>max_majors):
+            nsubint += 1
+        pmax = int(math.ceil(pmax/float(nsubint)))*nsubint
+        pmin = int(math.floor(pmin/float(nsubint)))*nsubint
+        assert (pmax-pmin) % nsubint == 0
+        ymin = basey**pmin
+        ymax = basey**pmax
+
+        major_yticks = tuple(basey**i for i in xrange(pmin, pmax+1, nsubint))
+        minor_yticks = tuple(basey**i for i in xrange(pmin, pmax+1, 1) if i%nsubint!=0)
+
+        ax = self.axes[0]
+        ax.set_ylim(ymin, ymax)
+        ax.yaxis.set_ticks(major_yticks)
+        ax.yaxis.set_ticks(minor_yticks, minor=True)
+        ax.set_title(u'{}\niteration={}, t={}'.format(self.fig_title, iteration, time))
+        ax.relim()
+    
+    def _draw(self):
+        if (not self.has_gui_running):
+            return
+        self.fig.canvas.draw()
+        self.fig.show()
+        plt.pause(0.001)
+
+    def _savefig(self, iteration):
+        filename = self.filename.format(ite='{:05}'.format(iteration))
+        self.fig.savefig(filename, dpi=self.fig.dpi, bbox_inches='tight')
+    
+    def _on_close(self, event):
+        self.has_gui_running = False
+
+    def _on_key_press(self, event):
+        key = event.key
+        if key == 'q':
+            plt.close(self.fig)
+            self.has_gui_running = False
+    
 
 if __name__ == '__main__':
     from hysop.tools.sympy_utils import round_expr