From da58b5f5940add17ec2fae8b680dc29a1886ae66 Mon Sep 17 00:00:00 2001
From: Jean-Baptiste Keck <Jean-Baptiste.Keck@imag.fr>
Date: Thu, 21 Feb 2019 16:59:03 +0100
Subject: [PATCH] compute energy host kernels
---
examples/taylor_green/taylor_green.py | 2 +-
hysop/numerics/fft/fft.py | 4 +-
hysop/numerics/fft/host_fft.py | 137 ++++++++++++++++++-----
hysop/operator/base/spectral_operator.py | 14 ++-
hysop/tools/numba_utils.py | 8 +-
5 files changed, 130 insertions(+), 35 deletions(-)
diff --git a/examples/taylor_green/taylor_green.py b/examples/taylor_green/taylor_green.py
index 24cb866ae..12196ca9b 100644
--- a/examples/taylor_green/taylor_green.py
+++ b/examples/taylor_green/taylor_green.py
@@ -42,7 +42,7 @@ def compute(args):
from hysop.operators import DirectionalAdvection, DirectionalStretchingDiffusion, \
DirectionalDiffusion, DirectionalStretching, \
StaticDirectionalStretching, Diffusion, \
- PoissonCurl, AdaptiveTimeStep, \
+ PoissonCurl, AdaptiveTimeStep, \
Enstrophy, MinMaxFieldStatistics, StrangSplitting, \
ParameterPlotter, Advection, MinMaxGradientStatistics
diff --git a/hysop/numerics/fft/fft.py b/hysop/numerics/fft/fft.py
index 1e78f21c1..0b924aaaa 100644
--- a/hysop/numerics/fft/fft.py
+++ b/hysop/numerics/fft/fft.py
@@ -713,12 +713,12 @@ class FFTI(object):
"""Plan to compute energy from src to energy."""
assert src.ndim == len(transforms)
assert dst.ndim == 1
- N = src.shape
+ N = tuple(int(_) for _ in src.shape)
K2 = ()
NS2 = ()
C2C = ()
for (Ni, Ti) in zip(N, transforms):
- c2c = STU.is_C2C(Ti)
+ c2c = int(STU.is_C2C(Ti))
Ki = Ni//2 if c2c else Ni-1
K2 += (Ki**2,)
NS2 += ((Ni+1)//2,)
diff --git a/hysop/numerics/fft/host_fft.py b/hysop/numerics/fft/host_fft.py
index 06aa1d620..321043944 100644
--- a/hysop/numerics/fft/host_fft.py
+++ b/hysop/numerics/fft/host_fft.py
@@ -6,11 +6,14 @@ OpenCl backend base interface for fast Fourier Transforms.
:class:`~hysop.numerics.host_fft.HostFFTQueue`
"""
-import psutil
+import psutil, functools
import numpy as np
+import numba as nb
-from hysop import __FFTW_NUM_THREADS__, __FFTW_PLANNER_EFFORT__, __FFTW_PLANNER_TIMELIMIT__
+from hysop import __FFTW_NUM_THREADS__, __FFTW_PLANNER_EFFORT__, __FFTW_PLANNER_TIMELIMIT__, \
+ __DEFAULT_NUMBA_TARGET__
from hysop.tools.types import first_not_None, check_instance
+from hysop.tools.numba_utils import make_numba_signature
from hysop.backend.host.host_array_backend import HostArrayBackend
from hysop.backend.host.host_array import HostArray
from hysop.numerics.fft.fft import FFTQueueI, FFTPlanI, FFTI
@@ -115,30 +118,114 @@ class HostFFTI(FFTI):
buf[slc] = 0
return exec_fill_zeros
- def plan_compute_energy(self, tg, src, dst, transforms, method='round'):
- """Plan to compute energy from src to energy."""
+ def plan_compute_energy(self, tg, src, dst, transforms,
+ method='round', target=None):
+ """
+ Plan to compute energy from src array to dst array using given transforms,
+ method (round or weighted) and numba target.
+ """
(N, NS2, C2C) = super(HostFFTI, self).plan_compute_energy(tg, src, dst, transforms)
- def exec_compute_energy(src=src, dst=dst, C2C=C2C, NS2=NS2):
- dst[...] = 0.0
- for idx in np.ndindex(*src.shape):
- Ek = (abs(src[idx])**2) / 2.0
- K = np.asarray(idx, dtype=np.int64)
- for (i, (Ki, Ni, ns2, c2c)) in enumerate(zip(K, N, NS2, C2C)):
- if c2c and Ki>ns2:
- K[i] = Ni - Ki
- k = (K.dot(K))**0.5
- if method=='round':
- i = int(round(k))
- dst[i] += Ek
- elif method=='weighted':
- i = int(k)
- a = k-i
- dst[i] += (1-a)*Ek
- dst[min(i+1,dst.size-1)] += a*Ek
- else:
- msg='Unknown method {}'.format(method)
- raise NotImplementedError(msg)
- print dst
+
+ target = first_not_None(target, __DEFAULT_NUMBA_TARGET__)
+ args = (src,)+N+NS2+C2C+(dst,)
+ signature, layout = make_numba_signature(*args)
+
+ dim = src.ndim
+ if (dim == 1):
+ @nb.guvectorize([signature], layout, target=target,
+ nopython=True, cache=True)
+ def compute_energy_1d(src, Nx, NS2x, C2Cx, dst):
+ for _ in range(dst.size):
+ dst[_] = 0.0
+ for ix in range(Nx):
+ kx = ((Nx-ix) if (C2Cx and ix>NS2x) else ix)
+ Xi = src[ix]
+ Ei = 0.5*(abs(Xi)**2)
+ dst[kx] += Ei
+ exec_compute_energy = functools.partial(compute_energy_1d, *args)
+ elif (dim == 2):
+ if (method=='round'):
+ @nb.guvectorize([signature], layout, target=target,
+ nopython=True, cache=True)
+ def compute_energy_2d_round(src, Ny, Nx, NS2y, NS2x, C2Cy, C2Cx, dst):
+ for _ in range(dst.size):
+ dst[_] = 0.0
+ for iy in range(Ny):
+ ky = ((Ny-iy) if (C2Cy and iy>NS2y) else iy)
+ for ix in range(Nx):
+ kx = ((Nx-ix) if (C2Cx and ix>NS2x) else ix)
+ Xi = src[iy,ix]
+ Ei = 0.5*(abs(Xi)**2)
+ k = int(round((kx**2 + ky**2)**0.5))
+ dst[k] += Ei
+ exec_compute_energy = functools.partial(compute_energy_2d_round, *args)
+ elif method=='weighted':
+ @nb.guvectorize([signature], layout, target=target,
+ nopython=True, cache=True)
+ def compute_energy_2d_weighted(src, Ny, Nx, NS2y, NS2x, C2Cy, C2Cx, dst):
+ for _ in range(dst.size):
+ dst[_] = 0.0
+ for iy in range(Ny):
+ ky = ((Ny-iy) if (C2Cy and iy>NS2y) else iy)
+ for ix in range(Nx):
+ kx = ((Nx-ix) if (C2Cx and ix>NS2x) else ix)
+ Xi = src[iy,ix]
+ Ei = 0.5*(abs(Xi)**2)
+ k = (kx**2 + ky**2)**0.5
+ i = int(k)
+ j = min(i+1, dst.size-1)
+ a = k-i
+ dst[i] += (1-a)*Ei
+ dst[j] += a *Ei
+ exec_compute_energy = functools.partial(compute_energy_2d_weighted, *args)
+ else:
+ msg='Unknown method {}.'.format(method)
+ raise ValueError(msg)
+ elif (dim == 3):
+ if (method=='round'):
+ @nb.guvectorize([signature], layout, target=target,
+ nopython=True, cache=True)
+ def compute_energy_3d_round(src, Nz, Ny, Nx, NS2z, NS2y, NS2x, C2Cz, C2Cy, C2Cx, dst):
+ for _ in range(dst.size):
+ dst[_] = 0.0
+ for iz in range(Nz):
+ kz = ((Nz-iz) if (C2Cz and iz>NS2z) else iz)
+ for iy in range(Ny):
+ ky = ((Ny-iy) if (C2Cy and iy>NS2y) else iy)
+ for ix in range(Nx):
+ kx = ((Nx-ix) if (C2Cx and ix>NS2x) else ix)
+ Xi = src[iy,ix,iz]
+ Ei = 0.5*(abs(Xi)**2)
+ k = int(round((kx**2 + ky**2 + kz**2)**0.5))
+ dst[k] += Ei
+ exec_compute_energy = functools.partial(compute_energy_3d_round, *args)
+ elif method=='weighted':
+ @nb.guvectorize([signature], layout, target=target,
+ nopython=True, cache=True)
+ def compute_energy_3d_weighted(src, Nz, Ny, Nx, NS2z, NS2y, NS2x, C2Cz, C2Cy, C2Cx, dst):
+ for _ in range(dst.size):
+ dst[_] = 0.0
+ for iz in range(Nz):
+ kz = ((Nz-iz) if (C2Cz and iz>NS2z) else iz)
+ for iy in range(Ny):
+ ky = ((Ny-iy) if (C2Cy and iy>NS2y) else iy)
+ for ix in range(Nx):
+ kx = ((Nx-ix) if (C2Cx and ix>NS2x) else ix)
+ Xi = src[iy,ix,iz]
+ Ei = 0.5*(abs(Xi)**2)
+ k = (kx**2 + ky**2 + kz**2)**0.5
+ i = int(k)
+ j = min(i+1, dst.size-1)
+ a = k-i
+ dst[i] += (1-a)*Ei
+ dst[j] += a *Ei
+ exec_compute_energy = functools.partial(compute_energy_3d_weighted, *args)
+ else:
+ msg='Unknown method {}.'.format(method)
+ raise ValueError(msg)
+ else:
+ msg='{}D compute_energy has not been implemented yet.'
+ raise NotImplementedError(msg.format(dim))
return exec_compute_energy
diff --git a/hysop/operator/base/spectral_operator.py b/hysop/operator/base/spectral_operator.py
index 12ffa0bea..1ede8be34 100644
--- a/hysop/operator/base/spectral_operator.py
+++ b/hysop/operator/base/spectral_operator.py
@@ -1884,6 +1884,8 @@ SPECTRAL TRANSFORM SETUP
compute_energy_queue += FFTI.plan_compute_energy(tg=tg,
src=spectral_buffer, dst=energy_buffer,
transforms=self.permuted_transforms)
+ if self._do_plot_energy:
+ self._energy_plotter = EnergyPlot(ioparams=self._plot_energy_ioparams, energy_buffer=energy_buffer)
else:
compute_energy_queue = None
@@ -1936,8 +1938,14 @@ SPECTRAL TRANSFORM SETUP
should_plot = (frequency>0) and (ite % frequency == 0)
should_plot |= simulation.is_time_of_interest
if should_plot:
- return self._plot_energy()
+ filename = self._plot_energy_ioparams.filename
+ filename = filename.replace('{fname}', self.field.name).replace('{it}', ite)
+ evt = self._energy_plotter.update()
+ self._energy_plotter.filename(filename)
+ return evt
- def _plot_energy(self):
- pass
+
+class EnergyPlot(object):
+ def __init__(self, energy_buffer):
+ check_instance(energy_buffer, npw.ndarray)
diff --git a/hysop/tools/numba_utils.py b/hysop/tools/numba_utils.py
index e8abd9cd4..069ed7caa 100644
--- a/hysop/tools/numba_utils.py
+++ b/hysop/tools/numba_utils.py
@@ -9,9 +9,9 @@ def make_numba_signature(*args, **kwds):
msg='Unknown kwds {}.'.forma(kwds.keys())
raise RuntimeError(kwds)
dtype_to_ntype = {
- int: np.int32,
- long: np.int64,
- float: np.float64,
+ int: nb.int32,
+ long: nb.int64,
+ float: nb.float64,
np.int8: nb.int8,
np.int16: nb.int16,
@@ -78,7 +78,7 @@ def make_numba_signature(*args, **kwds):
na = nb.types.Array(dtype=dtype, ndim=a.ndim, layout='A')
numba_layout += (format_shape(*a.shape),)
elif isinstance(a, np.dtype):
- na = dtype_to_ntype[a.type]
+ a = dtype_to_ntype[a.type]
numba_layout += ('()',)
elif isinstance(a, type):
na = dtype_to_ntype[a]
--
GitLab