diff --git a/examples/taylor_green/taylor_green.py b/examples/taylor_green/taylor_green.py index 6dc1100c884e8dbc650363b9b9930b9e210b811e..725e57758bb588639988581115bea7bfe9f3736c 100644 --- a/examples/taylor_green/taylor_green.py +++ b/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) diff --git a/hysop/backend/device/opencl/operator/poisson_curl.py b/hysop/backend/device/opencl/operator/poisson_curl.py index db26cd672fa54d5e914a6222c5b3291c993f9894..4f98a682eeefd41ec874c0c49669d6bd551c2f98 100644 --- a/hysop/backend/device/opencl/operator/poisson_curl.py +++ b/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) diff --git a/hysop/backend/host/python/operator/poisson_curl.py b/hysop/backend/host/python/operator/poisson_curl.py index 498cda5145ed4a01eb7bdd1b5bec9811a2358ebc..13e80586673b405bb3c7f797ba6042fffbd64ea9 100644 --- a/hysop/backend/host/python/operator/poisson_curl.py +++ b/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) diff --git a/hysop/operator/base/poisson_curl.py b/hysop/operator/base/poisson_curl.py index c37ff4ae974aa8eb4b2bb68339e9c91451d3eb1f..027f24d0e45ed2b7643e0436a427d17f37fc4636 100644 --- a/hysop/operator/base/poisson_curl.py +++ b/hysop/operator/base/poisson_curl.py @@ -9,7 +9,7 @@ from hysop.tools.io_utils import IOParams from hysop.fields.continuous_field import Field from hysop.parameters.scalar_parameter import ScalarParameter from hysop.parameters.buffer_parameter import BufferParameter -from hysop.operator.base.spectral_operator import SpectralOperatorBase, EnergyPlotter +from hysop.operator.base.spectral_operator import SpectralOperatorBase, EnergyPlotter, EnergyDumper from hysop.topology.cartesian_descriptor import CartesianTopologyDescriptors from hysop.fields.continuous_field import Field from hysop.symbolic.field import laplacian, grad @@ -24,7 +24,8 @@ class PoissonCurlOperatorBase(object): @debug def __init__(self, vorticity, velocity, variables, diffusion=None, dt=None, projection=None, - plot_velocity_energy=None, plot_vorticity_energy=None, + dump_energy=None, dump_velocity_energy=None, dump_input_vorticity_energy=None, dump_output_vorticity_energy=None, + plot_energy=None, plot_velocity_energy=None, plot_input_vorticity_energy=None, plot_output_vorticity_energy=None, plot_inout_vorticity_energy=None, **kwds): """ PoissonCurl operator to solve incompressible flows using various fft backends. @@ -48,15 +49,32 @@ class PoissonCurlOperatorBase(object): When active, projection is done prior to every solve, unless projection is an integer in which case it is done every given steps. This parameter is ignored for 2D fields and defaults to no projection. - plot_input_energy: IOParams, optional, defaults to None - Plot input field energy in a custom file. Defaults to no plot. - Energy is computed after filtering. - plot_output_energy: IOParams, optional, defaults to None - Plot output field energy in a custom file. Defaults to no plot. - Energy is computed after filtering. - If None is passed, no plots are generated. + dump_energy: IOParams, optional, defaults to None + Will set the default io parameter for all energy plotters. + dump_velocity_energy: IOParams, optional, defaults to None + Dump velocity field energy to a custom file. Defaults to no dump. + dump_input_vorticity_energy: IOParams, optional, defaults to None + Dump input vorticity field energy to a custom file. Defaults to no dump. + dump_output_vorticity_energy: IOParams, optional, defaults to None + Dump output vorticity field energy to a custom file. Defaults to no dump. + plot_energy: IOParams, optional, defaults to None + Will set the default io parameter for all energy plotters. + plot_velocity_energy: IOParams, optional, defaults to None + Plot velocity field energy and save the plot to a custom file. Defaults to no plot. + plot_input_vorticity_energy: IOParams, optional, defaults to None + Plot input vorticity field energy and save the plot to a custom file. Defaults to no plot. + plot_output_vorticity_energy: IOParams, optional, defaults to None + Plot output vorticity field energy and save the plot to a custom file. Defaults to no plot. + plot_inout_vorticity_energy: IOParams, optional, defaults to None + Plot vorticity field energy before and after diffusion and projection on the same graph. kwds : Base class parameters. + + Notes: + ------ + All dump energy arguments also enables scalar energy dumping. + This is not true for energy plotting. + Passing an integer instead of a IOParams will disable dump and plot arguments. """ check_instance(velocity, Field) @@ -65,8 +83,15 @@ class PoissonCurlOperatorBase(object): check_instance(diffusion, ScalarParameter, allow_none=True) check_instance(dt, ScalarParameter, allow_none=True) check_instance(projection, (FieldProjection, int), allow_none=True) - check_instance(plot_velocity_energy, IOParams, allow_none=True) - check_instance(plot_vorticity_energy, IOParams, allow_none=True) + check_instance(dump_energy, (IOParams,int), allow_none=True) + check_instance(dump_velocity_energy, (IOParams,int), allow_none=True) + check_instance(dump_input_vorticity_energy, (IOParams,int), allow_none=True) + check_instance(dump_output_vorticity_energy, (IOParams,int), allow_none=True) + check_instance(plot_energy, (IOParams,int), allow_none=True) + check_instance(plot_velocity_energy, (IOParams,int), allow_none=True) + check_instance(plot_input_vorticity_energy, (IOParams,int), allow_none=True) + check_instance(plot_output_vorticity_energy, (IOParams,int), allow_none=True) + check_instance(plot_inout_vorticity_energy, (IOParams,int), allow_none=True) assert velocity.domain is vorticity.domain, 'only one domain is supported' assert variables[velocity] is variables[vorticity], 'only one topology is supported' @@ -116,7 +141,6 @@ class PoissonCurlOperatorBase(object): vtopology = variables[velocity] wtopology = variables[vorticity] input_params = {} - output_params = {} input_fields = { vorticity: wtopology } output_fields = { velocity: vtopology } if should_diffuse: @@ -125,33 +149,37 @@ class PoissonCurlOperatorBase(object): input_params[dt.name] = dt if (should_diffuse or should_project): output_fields[vorticity] = wtopology - - compute_bwd_W = (should_diffuse or should_project) - do_compute_velocity_energy = (plot_velocity_energy is not None) - do_compute_vorticity_energy = (plot_vorticity_energy is not None) - do_compute_fwd_W_energy = do_compute_vorticity_energy and (not compute_bwd_W) - do_compute_bwd_W_energy = do_compute_vorticity_energy and compute_bwd_W - do_compute_bwd_U_energy = do_compute_velocity_energy - plot_velocity_energy_ioparams = plot_velocity_energy - plot_vorticity_energy_ioparams = plot_vorticity_energy - if do_compute_velocity_energy: - pname = 'E_{}'.format(velocity.name) - pename = 'E_{}'.format(velocity.pretty_name) - velocity_energy_parameter = BufferParameter( - name=pname, pretty_name=pename, - shape=None, dtype=None, initial_value=None) - output_params[velocity_energy_parameter.name] = velocity_energy_parameter - else: - velocity_energy_parameter = None - if do_compute_vorticity_energy: - pname = 'E_{}'.format(vorticity.name) - pename = 'E_{}'.format(vorticity.pretty_name) - vorticity_energy_parameter = BufferParameter( - name=pname, pretty_name=pename, - shape=None, dtype=None, initial_value=None) - output_params[vorticity_energy_parameter.name] = vorticity_energy_parameter - else: - vorticity_energy_parameter = None + + dump_Uout_E = first_not_None(dump_velocity_energy, dump_energy) + dump_Win_E = first_not_None(dump_input_vorticity_energy, dump_energy) + dump_Wout_E = first_not_None(dump_output_vorticity_energy, dump_energy) + + plot_Uout_E = first_not_None(plot_velocity_energy, plot_energy) + plot_Win_E = first_not_None(plot_input_vorticity_energy, plot_energy) + plot_Wout_E = first_not_None(plot_output_vorticity_energy, plot_energy) + plot_Winout_E = first_not_None(plot_inout_vorticity_energy, plot_energy) + + do_dump_Uout_E = isinstance(dump_Uout_E, IOParams) and (dump_Uout_E.frequency>=0) + do_dump_Win_E = isinstance(dump_Win_E, IOParams) and (dump_Win_E.frequency>=0) + do_dump_Wout_E = isinstance(dump_Wout_E, IOParams) and (dump_Wout_E.frequency>=0) + + do_plot_Uout_E = isinstance(plot_Uout_E, IOParams) and (plot_Uout_E.frequency>=0) + do_plot_Win_E = isinstance(plot_Win_E, IOParams) and (plot_Win_E.frequency>=0) + do_plot_Wout_E = isinstance(plot_Wout_E, IOParams) and (plot_Wout_E.frequency>=0) + do_plot_Winout_E = isinstance(plot_Winout_E, IOParams) and (plot_Winout_E.frequency>=0) + + do_compute_Uout_E, compute_Uout_E_freqs = EnergyDumper.do_compute_energy(dump_Uout_E, plot_Uout_E) + do_compute_Win_E, compute_Win_E_freqs = EnergyDumper.do_compute_energy(dump_Win_E, plot_Win_E, plot_Winout_E) + do_compute_Wout_E, compute_Wout_E_freqs = EnergyDumper.do_compute_energy(dump_Wout_E, plot_Wout_E, plot_Winout_E) + + if do_compute_Wout_E: + msg='Cannot compute output vorticity energy because there is no output vorticity !' + assert (should_diffuse or should_project), msg + + output_params = {} + compute_Win_E_param = EnergyDumper.build_energy_parameter(do_compute=do_compute_Win_E, field=vorticity, output_params=output_params, prefix='in') + compute_Uout_E_param = EnergyDumper.build_energy_parameter(do_compute=do_compute_Uout_E, field=velocity, output_params=output_params, prefix='out') + compute_Wout_E_param = EnergyDumper.build_energy_parameter(do_compute=do_compute_Wout_E, field=vorticity, output_params=output_params, prefix='out') super(PoissonCurlOperatorBase, self).__init__(input_fields=input_fields, output_fields=output_fields, input_params=input_params, @@ -168,18 +196,15 @@ class PoissonCurlOperatorBase(object): self.should_project = should_project self.projection = projection self.do_project = do_project - self.compute_bwd_W = (should_diffuse or should_project) - - self.do_compute_velocity_energy = do_compute_velocity_energy - self.do_compute_vorticity_energy = do_compute_vorticity_energy - self.do_compute_fwd_W_energy = do_compute_fwd_W_energy - self.do_compute_bwd_W_energy = do_compute_bwd_W_energy - self.do_compute_bwd_U_energy = do_compute_bwd_U_energy - self.plot_velocity_energy_ioparams = plot_velocity_energy_ioparams - self.plot_vorticity_energy_ioparams = plot_vorticity_energy_ioparams - self.velocity_energy_parameter = velocity_energy_parameter - self.vorticity_energy_parameter = vorticity_energy_parameter - + + self.do_compute_energy = {'Uout': do_compute_Uout_E, 'Win': do_compute_Win_E, 'Wout': do_compute_Wout_E} + self.do_dump_energy = {'Uout': do_dump_Uout_E, 'Win': do_dump_Win_E, 'Wout': do_dump_Wout_E} + self.do_plot_energy = {'Uout': do_plot_Uout_E, 'Win': do_plot_Win_E, 'Wout': do_plot_Wout_E, 'Winout': do_plot_Winout_E} + self.dump_energy_ioparams = {'Uout': dump_Uout_E, 'Win': dump_Win_E, 'Wout': dump_Wout_E} + self.plot_energy_ioparams = {'Uout': plot_Uout_E, 'Win': plot_Win_E, 'Wout': plot_Wout_E, 'Winout': plot_Winout_E} + self.compute_energy_parameters = {'Uout': compute_Uout_E_param, 'Win': compute_Win_E_param, 'Wout': compute_Wout_E_param } + self.compute_energy_frequencies = {'Uout': compute_Uout_E_freqs, 'Win': compute_Win_E_freqs, 'Wout': compute_Wout_E_freqs } + @debug def discretize(self): @@ -202,25 +227,22 @@ class SpectralPoissonCurlOperatorBase(PoissonCurlOperatorBase, SpectralOperatorB dim = self.dim should_diffuse, should_project = self.should_diffuse, self.should_project - U_energy_ioparams = self.plot_velocity_energy_ioparams - W_energy_ioparams = self.plot_vorticity_energy_ioparams - do_compute_fwd_W_energy = self.do_compute_fwd_W_energy - do_compute_bwd_W_energy = self.do_compute_bwd_W_energy - do_compute_bwd_U_energy = self.do_compute_bwd_U_energy - compute_fwd_W_energy = None if (not do_compute_fwd_W_energy) else W_energy_ioparams.frequency - compute_bwd_W_energy = None if (not do_compute_bwd_W_energy) else W_energy_ioparams.frequency - compute_bwd_U_energy = None if (not do_compute_bwd_U_energy) else U_energy_ioparams.frequency + de_iops = self.dump_energy_ioparams + ce_freqs = self.compute_energy_frequencies # build spectral transforms tg = self.new_transform_group() W_forward_transforms = to_tuple(tg.require_forward_transform(vorticity, - compute_energy=compute_fwd_W_energy)) + compute_energy_frequencies=ce_freqs['Win'], + dump_energy=de_iops['Win'])) U_backward_transforms = to_tuple(tg.require_backward_transform(velocity, custom_input_buffer='B0', - compute_energy=compute_bwd_U_energy)) + compute_energy_frequencies=ce_freqs['Uout'], + dump_energy=de_iops['Uout'])) if (should_diffuse or should_project): W_backward_transforms = to_tuple(tg.require_backward_transform(vorticity, - compute_energy=compute_bwd_W_energy)) + compute_energy_frequencies=ce_freqs['Wout'], + dump_energy=de_iops['Wout'])) else: assert compute_bwd_W_energy == 0 W_backward_transforms = (None,)*dim @@ -288,69 +310,89 @@ class SpectralPoissonCurlOperatorBase(PoissonCurlOperatorBase, SpectralOperatorB if self.discretized: return super(SpectralPoissonCurlOperatorBase, self).discretize() + + fig_titles = {'Win': u'Energy of input vorticity {}', + 'Uout': u'Energy of output velocity {}', + 'Wout': u'Energy of output vorticity {}' } + get_transforms = { 'Win': self.W_forward_transforms, + 'Wout': self.W_backward_transforms, + 'Uout': self.U_backward_transforms } + get_dfield = { 'Win': self.dW, + 'Wout': self.dW, + 'Uout': self.dU } - if self.do_compute_velocity_energy: - Ep = self.velocity_energy_parameter - iop = self.plot_velocity_energy_ioparams - E_params = tuple(U_bt._energy_parameter for U_bt in self.U_backward_transforms) - assert all(p.size is not None for p in E_params) + compute_fns = () + for (k, v) in self.do_compute_energy.iteritems(): + if not v: + assert not self.do_dump_energy[k] + assert not self.do_plot_energy[k] + continue + dfield = get_dfield[k] + for tr in get_transforms[k]: + if tr._energy_parameter is None: + msg='Energy parameter of {}.{} has not been build.' + raise RuntimeError(msg.format(tr.field.name, + 'forward' if tr.is_forward else 'backward')) + E_params = tuple(tr._energy_parameter for tr in get_transforms[k]) E_buffers = tuple(p.value for p in E_params) - E_size = max(p.size for p in E_params) + E_size = max(p.size for p in E_params) + + Ep = self.compute_energy_parameters[k] Ep.reallocate_buffer(shape=(E_size,), dtype=self.dU.dtype) - pname = u'{}.{}'.format(self.__class__.__name__, self.dU.pretty_name.decode('utf-8')) - energy_parameters = { pname : Ep } - filename = self.plot_velocity_energy_ioparams.filename.format(fname=self.dU.name, ite='') + '.txt' - plt = EnergyPlotter(energy_parameters=energy_parameters, - fig_title=u'Energy of velocity {}'.format(self.U.pretty_name.decode('utf-8')), - save_txt=True, filename=filename) - def compute_velocity_energy_and_plot(simulation, iop=iop, - plt=plt, dst=Ep._value, srcs=E_buffers): - if simulation.should_dump(frequency=iop.frequency, with_last=True): - dst[...] = 0.0 - for src in srcs: - dst[:src.size] += src - filename = iop.filename - filename = filename.replace('{fname}', 'U') - filename = filename.replace('{ite}', '{:05}'.format(simulation.current_iteration)) - plt.update(iteration=simulation.current_iteration, time=simulation.t()) - plt.save(filename) - self.compute_velocity_energy_and_plot = compute_velocity_energy_and_plot - else: - self.compute_velocity_energy_and_plot = None - - if self.do_compute_vorticity_energy: - Ep = self.vorticity_energy_parameter - iop = self.plot_vorticity_energy_ioparams - if self.do_compute_bwd_W_energy: - E_params = tuple(W_bt._energy_parameter for W_bt in self.W_backward_transforms) - elif self.do_compute_fwd_W_energy: - E_params = tuple(W_ft._energy_parameter for W_ft in self.W_forward_transforms) + + if self.do_dump_energy[k]: + iop = self.dump_energy_ioparams[k] + assert (iop is not None) + dmp = EnergyDumper(energy_parameter=Ep, io_params=iop, fname=k) else: - raise RuntimeError - assert all(p.size is not None for p in E_params) - E_buffers = tuple(p.value for p in E_params) - E_size = max(p.size for p in E_params) - Ep.reallocate_buffer(shape=(E_size,), dtype=self.dW.dtype) - filename = self.plot_vorticity_energy_ioparams.filename.format(fname=self.dW.name, ite='') + '.txt' - pname = u'{}.{}'.format(self.__class__.__name__, self.dW.pretty_name.decode('utf-8')) - energy_parameters = { pname : Ep } - plt = EnergyPlotter(energy_parameters=energy_parameters, - fig_title=u'Energy of vorticity {}'.format(self.W.pretty_name.decode('utf-8')), - save_txt=True, filename=filename) - def compute_vorticity_energy_and_plot(simulation, iop=iop, - plt=plt, dst=Ep._value, srcs=E_buffers): - if simulation.should_dump(frequency=iop.frequency, with_last=True): + dmp = None + + if self.do_plot_energy[k]: + iop = self.plot_energy_ioparams[k] + assert (iop is not None) + pname = u'{}.{}'.format(self.__class__.__name__, dfield.pretty_name.decode('utf-8')) + energy_parameters = { pname : Ep } + plt = EnergyPlotter(energy_parameters=energy_parameters, + io_params=iop, fname=k, + fig_title=fig_titles[k].format(dfield.pretty_name.decode('utf-8'))) + else: + plt = None + + freqs = self.compute_energy_frequencies[k] + + def compute_fn(simulation, plt=plt, dmp=dmp, dst=Ep._value, srcs=E_buffers, freqs=freqs): + should_compute = any(simulation.should_dump(frequency=f, with_last=True) for f in freqs) + if should_compute: dst[...] = 0.0 for src in srcs: dst[:src.size] += src - filename = iop.filename - filename = filename.replace('{fname}', 'W') - filename = filename.replace('{ite}', '{:05}'.format(simulation.current_iteration)) - plt.update(iteration=simulation.current_iteration, time=simulation.t()) - plt.save(filename) - self.compute_vorticity_energy_and_plot = compute_vorticity_energy_and_plot - else: - self.compute_vorticity_energy_and_plot = None + if (dmp is not None): + dmp.update(simulation=simulation, wait_for=None) + if (plt is not None): + plt.update(simulation=simulation, wait_for=None) + + compute_fns += (compute_fn,) + + assert 'Winout' not in self.do_compute_energy + assert 'Winout' not in self.do_dump_energy + if self.do_plot_energy['Winout']: + iop = self.plot_energy_ioparams['Winout'] + assert (iop is not None) + pname_in = u'{}.input.{}'.format(self.__class__.__name__, self.dW.pretty_name.decode('utf-8')) + pname_out = u'{}.output.{}'.format(self.__class__.__name__, self.dW.pretty_name.decode('utf-8')) + energy_parameters = { pname_in: self.compute_energy_parameters['Win'], + pname_out: self.compute_energy_parameters['Wout'] } + plt = EnergyPlotter(energy_parameters=energy_parameters, + io_params=iop, fname='Winout', + fig_title='Input and output vorticity energy (diffusion={:.2e}, project={})'.format(self.nu(), self.projection)) + def compute_fn(simulation, plt=plt): + plt.update(simulation=simulation, wait_for=None) + compute_fns += (compute_fn,) + + def update_energy(simulation, compute_fns=compute_fns): + for fn in compute_fns: + fn(simulation=simulation) + self.update_energy = update_energy kd1s, kd2s = self.kd1s, self.kd2s if self.should_project: diff --git a/hysop/operator/base/spectral_operator.py b/hysop/operator/base/spectral_operator.py index 2c276c880615650ec458db173597b867762116a5..5b5779eebf97228147854d355b78325ee09b0f0a 100644 --- a/hysop/operator/base/spectral_operator.py +++ b/hysop/operator/base/spectral_operator.py @@ -2,20 +2,18 @@ import warnings, math, os import sympy as sm import numpy as np -import matplotlib -import matplotlib.pyplot as plt from hysop.constants import BoundaryCondition, BoundaryExtension, TransformType, \ MemoryOrdering, TranspositionState, Backend, \ SpectralTransformAction, Implementation from hysop.tools.misc import compute_nbytes -from hysop.tools.types import check_instance, to_tuple, first_not_None +from hysop.tools.types import check_instance, to_tuple, first_not_None, to_set from hysop.tools.decorators import debug from hysop.tools.units import bytes2str from hysop.tools.numerics import is_fp, is_complex, complex_to_float_dtype, \ float_to_complex_dtype, determine_fp_types from hysop.tools.io_utils import IOParams -from hysop.tools.spectral_utils import SpectralTransformUtils as STU +from hysop.tools.spectral_utils import SpectralTransformUtils as STU, EnergyPlotter, EnergyDumper from hysop.core.mpi import main_rank from hysop.core.arrays.array_backend import ArrayBackend from hysop.core.arrays.array import Array @@ -562,7 +560,8 @@ class SpectralTransformGroup(object): @_not_initialized def require_forward_transform(self, field, axes=None, transform_tag=None, custom_output_buffer=None, action=None, - compute_energy=None, plot_energy=None): + dump_energy=None, plot_energy=None, + **kwds): """ Tells this SpectralTransformGroup to build a forward SpectralTransform on given field. Only specified axes are transformed. @@ -606,15 +605,34 @@ class SpectralTransformGroup(object): compute slices of the output buffer. SpectralTransformAction.ACCUMULATE will sum the current content of the buffer with the result of the forward transform. - compute_energy: int, optional, defaults to None + dump_energy: IOParams, optional, defaults to None Compute the energy for each wavenumber at given frequency after each transform. + If None is passed, no files are generated (default behaviour). plot_energy: IOParams, optional, defaults to None Plot field energy after each call to the forward transform to a custom file. - This option implies compute_energy == io_params.frequency. If None is passed, no plots are generated (default behaviour). + compute_energy_frequencies: array like of integers, optional, defaults to None + Extra frequencies where to compute energy. + Notes + ----- IOParams filename is formatted before being used: - {fname} is replaced with field name - {ite} is replaced with simulation iteration id + {fname} is replaced with discrete field name + {ite} is replaced with simulation iteration id for plotting and '' for file dumping. + + + dump_energy plot_energy result + None None nothing + iop0 0 energy is computed and dumped every iop0.frequency iterations + 0 iop1 energy is computed and dumped every iop1.frequency iterations + iop0 iop1 energy is computed every iop1.frequency and iop2.frequency iterations + dumped every iop0.frequency + plotted every iop1.frequency + + About frequency: + if (frequency<0) no dump + if (frequency==0) dump at time of interests and last iteration + if (frequency>=0) dump at time of interests, last iteration and every freq iterations + """ transform_tag = first_not_None(transform_tag, 'default') action = first_not_None(action, SpectralTransformAction.OVERWRITE) @@ -636,8 +654,9 @@ class SpectralTransformGroup(object): symbolic_transform=transforms[idx], custom_output_buffer=custom_output_buffer, action=action, - compute_energy=compute_energy, - plot_energy=plot_energy) + dump_energy=dump_energy, + plot_energy=plot_energy, + **kwds) self._forward_transforms[(f,axes,transform_tag)] = planned_transforms[idx] else: assert (field,axes,transform_tag) not in self._forward_transforms, msg.format(field.name, axes, transform_tag) @@ -649,8 +668,9 @@ class SpectralTransformGroup(object): symbolic_transform=transforms, custom_output_buffer=custom_output_buffer, action=action, - compute_energy=compute_energy, - plot_energy=plot_energy) + dump_energy=dump_energy, + plot_energy=plot_energy, + **kwds) self._forward_transforms[(field,axes,transform_tag)] = planned_transforms return planned_transforms @@ -659,7 +679,8 @@ class SpectralTransformGroup(object): custom_input_buffer=None, matching_forward_transform=None, action=None, - compute_energy=None, plot_energy=None): + dump_energy=None, plot_energy=None, + **kwds): """ Same as require_forward_transform but for backward transforms. This corresponds to the following backward transform mappings: @@ -716,15 +737,35 @@ class SpectralTransformGroup(object): compute slices of the given output field. SpectralTransformAction.ACCUMULATE will sum the current content of the field with the result of the backward transform. - compute_energy: int, optional, defaults to None - Compute the energy for each wavenumber at given frequency before each call to the transform. + dump_energy: IOParams, optional, defaults to None + Compute the energy for each wavenumber at given frequency before each transform. + If None is passed, no files are generated (default behaviour). plot_energy: IOParams, optional, defaults to None - Plot field energy before each call to the forward transform to a custom file. - This option implies compute_energy == io_params.frequency. + Plot field energy before each call to the backward transform to a custom file. If None is passed, no plots are generated (default behaviour). + compute_energy_frequencies: array like of integers, optional, defaults to None + Extra frequencies where to compute energy. + + Notes + ----- IOParams filename is formatted before being used: - {fname} is replaced with field name - {ite} is replaced with simulation iteration id + {fname} is replaced with discrete field name + {ite} is replaced with simulation iteration id for plotting and '' for file dumping. + + + dump_energy plot_energy result + None None nothing + iop0 0 energy is computed and dumped every iop0.frequency iterations + 0 iop1 energy is computed and dumped every iop1.frequency iterations + iop0 iop1 energy is computed every iop1.frequency and iop2.frequency iterations + dumped every iop0.frequency + plotted every iop1.frequency + + About frequency: + if (frequency<0) no dump + if (frequency==0) dump at time of interests and last iteration + if (frequency>=0) dump at time of interests, last iteration and every freq iterations + """ transform_tag = first_not_None(transform_tag, 'default') action = first_not_None(action, SpectralTransformAction.OVERWRITE) @@ -745,8 +786,9 @@ class SpectralTransformGroup(object): custom_input_buffer=custom_input_buffer, matching_forward_transform=matching_forward_transform, action=action, - compute_energy=compute_energy, - plot_energy=plot_energy) + dump_energy=dump_energy, + plot_energy=plot_energy, + **kwds) self._backward_transforms[(f,axes,transform_tag)] = planned_transforms[idx] else: assert (field,axes,transform_tag) not in self._backward_transforms, msg.format(field.name, axes, transform_tag) @@ -758,8 +800,9 @@ class SpectralTransformGroup(object): custom_input_buffer=custom_input_buffer, matching_forward_transform=matching_forward_transform, action=action, - compute_energy=compute_energy, - plot_energy=plot_energy) + dump_energy=dump_energy, + plot_energy=plot_energy, + **kwds) self._backward_transforms[(field,axes,transform_tag)] = planned_transforms return planned_transforms @@ -832,14 +875,14 @@ class PlannedSpectralTransform(object): def __init__(self, transform_group, tag, symbolic_transform, action, custom_input_buffer=None, custom_output_buffer=None, matching_forward_transform=None, - compute_energy=None, plot_energy=None): + dump_energy=None, plot_energy=None,compute_energy_frequencies=None): check_instance(transform_group, SpectralTransformGroup) check_instance(transform_group.op, SpectralOperatorBase) check_instance(tag, str) check_instance(symbolic_transform, AppliedSpectralTransform) check_instance(action, SpectralTransformAction) - check_instance(compute_energy, int, allow_none=True) + check_instance(dump_energy, IOParams, allow_none=True) check_instance(plot_energy, IOParams, allow_none=True) assert custom_input_buffer in (None, 'B0', 'B1', 'auto'), custom_input_buffer assert custom_output_buffer in (None, 'B0', 'B1', 'auto'), custom_output_buffer @@ -855,11 +898,22 @@ class PlannedSpectralTransform(object): self._custom_output_buffer = custom_output_buffer self._matching_forward_transform = matching_forward_transform self._action = action + + self._do_dump_energy = (dump_energy is not None) and (dump_energy.frequency>=0) + self._do_plot_energy = (plot_energy is not None) and (plot_energy.frequency>=0) - self._do_compute_energy = (compute_energy is not None) or (plot_energy is not None) - self._do_plot_energy = (plot_energy is not None) - self._compute_energy_frequency = first_not_None(getattr(plot_energy, 'frequency', None), compute_energy) + compute_energy_frequencies = to_set(first_not_None(compute_energy_frequencies, set())) + if self._do_dump_energy: + compute_energy_frequencies.add(dump_energy.frequency) + if self._do_plot_energy: + compute_energy_frequencies.add(plot_energy.frequency) + compute_energy_frequencies = set(filter(lambda f: f>=0, compute_energy_frequencies)) + do_compute_energy = (len(compute_energy_frequencies)>0) + + self._do_compute_energy = do_compute_energy + self._compute_energy_frequencies = compute_energy_frequencies self._plot_energy_ioparams = plot_energy + self._dump_energy_ioparams = dump_energy if self._do_compute_energy: ename = 'E{}_{}'.format('f' if is_forward else 'b', field.name) @@ -869,6 +923,7 @@ class PlannedSpectralTransform(object): shape=None, dtype=None, initial_value=None) else: self._energy_parameter = None + self._energy_dumper = None self._energy_plotter = None if is_forward: @@ -1197,15 +1252,30 @@ class PlannedSpectralTransform(object): self._max_wavenumber = max_wavenumber self._energy_nbytes = energy_nbytes self._mutexes_nbytes = mutexes_nbytes - self._energy_parameter.reallocate_buffer(shape=(max_wavenumber+1,), dtype=dfield.dtype) + + Ep = self._energy_parameter + Ep.reallocate_buffer(shape=(max_wavenumber+1,), dtype=dfield.dtype) + + fname = fname='{}{}'.format(dfield.name, '_in' if is_forward else '_out') + + # build txt dumper + if self._do_dump_energy: + diop = self._dump_energy_ioparams + assert (diop is not None) + self._energy_dumper = EnergyDumper(energy_parameter=Ep, + io_params=self._dump_energy_ioparams, fname=fname) + + # build plotter if required if self._do_plot_energy: - filename = self._plot_energy_ioparams.filename.format(fname=dfield.name, ite='') + '.txt' + piop = self._plot_energy_ioparams + assert (piop is not None) pname = u'{}.{}.{}'.format(self.op.__class__.__name__, 'forward'if is_forward else 'backward', dfield.pretty_name.decode('utf-8')) energy_parameters = { pname: self._energy_parameter} self._energy_plotter = EnergyPlotter(energy_parameters=energy_parameters, - save_txt=True, filename=filename) + io_params=self._plot_energy_ioparams, + fname=fname) else: self._max_wavenumber = None self._energy_nbytes = None @@ -1938,7 +2008,7 @@ SPECTRAL TRANSFORM SETUP if self._do_compute_energy: field_buffer = self.input_buffer if self.is_forward else self.output_buffer spectral_buffer = self.output_buffer if self.is_forward else self.input_buffer - compute_energy_queue = FFTI.new_queue(tg=self, name='compute_energy') + compute_energy_queue = FFTI.new_queue(tg=self, name='dump_energy') compute_energy_queue += FFTI.plan_fill_zeros(tg=tg, a=energy_buffer, slices=(Ellipsis,)) if (mutexes_buffer is not None): unlock_mutexes = FFTI.plan_fill_zeros(tg=tg, a=mutexes_buffer, slices=(Ellipsis,)) @@ -1988,190 +2058,19 @@ SPECTRAL TRANSFORM SETUP def compute_energy(self, simulation, wait_for): msg='No simulation was passed in {}.__call__().'.format(type(self)) assert (simulation is not None), msg - frequency = self._compute_energy_frequency evt = wait_for - if simulation.should_dump(frequency=frequency, with_last=True): + should_compute_energy = any(simulation.should_dump(frequency=f, with_last=True) for f in self._compute_energy_frequencies) + if should_compute_energy: evt = self._compute_energy_queue(wait_for=evt) + if self._do_dump_energy: + self._energy_dumper.update(simulation=simulation, wait_for=evt) return evt def plot_energy(self, simulation, wait_for): msg='No simulation was passed in {}.__call__().'.format(type(self)) assert (simulation is not None), msg - frequency = self._plot_energy_ioparams.frequency - if simulation.should_dump(frequency=frequency, with_last=True): - wait_for.wait() - filename = self._plot_energy_ioparams.filename - filename = filename.replace('{fname}', self.field.name) - filename = filename.replace('{ite}', '{:05}'.format(simulation.current_iteration)) - self._energy_plotter.update(iteration=simulation.current_iteration, time=simulation.t()) - self._energy_plotter.save(filename) + evt = wait_for + self._energy_plotter.update(simulation=simulation, wait_for=evt) return wait_for -class EnergyPlotter(object): - def __init__(self, energy_parameters, fig_title=None, - axes_shape=(1,), figsize=(15,9), visu_rank=0, - basex=10, basey=10, - save_txt=True, filename=None, - **kwds): - super(EnergyPlotter, self).__init__(**kwds) - check_instance(axes_shape, tuple, minsize=1, allow_none=True) - check_instance(energy_parameters, dict, values=BufferParameter) - - should_draw = (visu_rank == main_rank) - - 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)[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 ** 2 / 1000 - if save_txt: - assert len(energy_parameters)==1, 'save_txt is not supported for multiple parameters.' - assert filename is not None - if os.path.isfile(filename): - os.remove(filename) - txt_formatter = {'float_kind': lambda x: '{:8.8f}'.format(x)} - f = open(filename, 'a') - header = u'# Evolution of the power spectrum of {}'.format(energy_parameters.keys()[0]) - f.write(header.encode('utf-8')) - f.write('\n# with first coefficient removed'.format(ulp)) - f.write('\n# ITERATION TIME log10(POWER_SPECTRUM[1:]))') - f.flush() - else: - f = None - txt_formatter = None - - self.fig = fig - self.axes = axes - self.lines = lines - - self.save_txt = save_txt - self.txt_formatter = txt_formatter - self.file = f - - self.running = running - self.should_draw = should_draw - self.basex = basex - self.basey = basey - - self.energy_parameters = energy_parameters - self.ulp = ulp - - def on_close(self, event): - self.running = False - - def on_key_press(self, event): - key = event.key - if key == 'q': - plt.close(self.fig) - self.running = False - - def draw(self): - if (not self.running): - return - if (not self.should_draw): - return - self.fig.canvas.draw() - self.fig.show() - plt.pause(0.001) - - def update(self, iteration, time, **kwds): - ymin=np.PINF - ymax=np.NINF - for p,line in zip(self.energy_parameters.values(), self.lines): - energy = p.value[1:].astype(dtype=np.float64) - - if self.save_txt: - assert (self.file is not None) - values = np.array2string(np.log10(energy), - formatter=self.txt_formatter, max_line_width=np.inf) - values = '\n{} {} {}'.format(iteration, time, values[1:-1]) - self.file.write(values) - self.file.flush() - - 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() - - self.draw() - - def save(self, filepath): - self.fig.savefig(filepath, dpi=self.fig.dpi, - bbox_inches='tight') diff --git a/hysop/simulation.py b/hysop/simulation.py index 5161c2790f9467c6a06db8d46e419ae197d79fc1..629a53904b229d90d315cbb7974ba1f7d00b4ebc 100644 --- a/hysop/simulation.py +++ b/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): diff --git a/hysop/tools/io_utils.py b/hysop/tools/io_utils.py index 997cf35d02333a3d33cbc04cf8014524186028d3..31d1406f29efa686748b214e2fbe15aaa80b7e04 100755 --- a/hysop/tools/io_utils.py +++ b/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): diff --git a/hysop/tools/spectral_utils.py b/hysop/tools/spectral_utils.py index 4f70d00ee508ef707f9113fff288cf679ea9bb3a..3f06e1b0e86fc1cad56d805e7f1104982a2f33eb 100644 --- a/hysop/tools/spectral_utils.py +++ b/hysop/tools/spectral_utils.py @@ -1,7 +1,11 @@ - +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