Add new 3D scalar field spectrum computation (not validated but seems ok)

HySoP/hysop/f2py/fftw2py.f90

@@ -307,5 +307,23 @@ contains
 
   end subroutine solve_curl_2d
 
+  !> Compute spectrum of a scalar field
+  !! @param[in] field
+  !! @param[out] spectrum
+  subroutine spectrum_3d(field, spectrum, wavelengths, ghosts, length)
+    real(pk),dimension(:,:,:),intent(in):: field
+    integer, dimension(3), intent(in) :: ghosts
+    real(pk),dimension(:), intent(inout) :: spectrum
+    real(pk),dimension(:), intent(inout) :: wavelengths
+    real(pk),intent(in) :: length
+    !f2py intent(in) :: field
+    !f2py intent(inout) :: spectrum
+    !f2py intent(inout) :: wavelengths
+
+    call r2c_3d_scal(field, ghosts)
+
+    call filter_spectrum_3d(spectrum, wavelengths, length)
+
+  end subroutine spectrum_3d
 
 end module fftw2py

HySoP/hysop/operator/discrete/spectrum.py

+# -*- coding: utf-8 -*-
+"""
+@file spectrum.py
+Discrete Spectrum operator using FFTW (fortran)
+"""
+try:
+    from hysop.f2py import fftw2py
+except ImportError:
+    from hysop.fakef2py import fftw2py
+from hysop.operator.discrete.discrete import DiscreteOperator
+from hysop.constants import debug, np, HYSOP_MPI_REAL
+from hysop.tools.profiler import profile
+from hysop.tools.io_utils import IO
+import hysop.tools.numpywrappers as npw
+from hysop.mpi import MPI
+import os
+
+
+class FFTSpectrum(DiscreteOperator):
+    """
+    Discretized Spectrum operator based on FFTW.
+
+    """
+    @debug
+    def __init__(self, field, prefix=None, **kwds):
+        """
+        Constructor.
+        @param[in] vorticity : field to compute.
+        """
+        # Discretisation of the input field
+        self.field = field
+
+        if self.field.nb_components > 1:
+            raise AttributeError("Vector case not yet implemented.")
+        # Base class initialisation
+        assert 'variables' not in kwds, 'variables parameter is useless.'
+        super(FFTSpectrum, self).__init__(variables=[field],
+                                          **kwds)
+        self.input = [self.field]
+        l = np.min(self.field.topology.mesh.discretization.resolution)
+        self._tmp = npw.zeros(((l - 1) / 2, ))
+        self._kx = npw.zeros(((l - 1) / 2, ))
+        self.res = npw.zeros(((l - 1) / 2, ))
+        self._prefix = "spectrum" if prefix is None else prefix
+        IO.check_dir(os.path.join(IO.default_path(),
+                                  self._prefix + "_00000.dat"))
+
+    @debug
+    @profile
+    def apply(self, simulation=None):
+        assert simulation is not None, \
+            "Missing dt value for diffusion computation."
+        ite = simulation.currentIteration
+        ghosts = self.field.topology.ghosts()
+        if self.field.dimension == 3:
+            fftw2py.spectrum_3d(self.field.data[0],
+                                self._tmp, self._kx,
+                                ghosts, np.min(self.domain.length))
+            if self.field.topology.size == 1:
+                self.res[...] = self._tmp
+            else:
+                self.field.topology.comm.Reduce(
+                    [self._tmp, self.res.shape[0], HYSOP_MPI_REAL],
+                    [self.res, self.res.shape[0], HYSOP_MPI_REAL],
+                    op=MPI.SUM, root=0)
+            if self.field.topology.rank == 0:
+                _file = open(os.path.join(
+                    IO.default_path(),
+                    self._prefix + "_{0:05d}.dat".format(ite)), 'w')
+                for i in xrange(self.res.shape[0]):
+                    _file.write("{0} {1}\n".format(
+                        self._kx[i], self.res[i]))
+                _file.close()
+        else:
+            raise ValueError("invalid problem dimension")
+
+    def finalize(self):
+        """
+        Clean memory (fftw plans and so on)
+        """
+        pass
+        # TODO : fix bug that occurs when several finalize
+        # of fft operators are called.
+        # fftw2py.clean_fftw_solver(self.vorticity.dimension)

HySoP/hysop/operator/spectrum.py

+"""
+@file spectrum.py
+"""
+from hysop.operator.computational import Computational
+from hysop.operator.discrete.spectrum import FFTSpectrum
+from hysop.constants import debug
+from hysop.operator.continuous import opsetup
+
+
+class Spectrum(Computational):
+    """
+    Fourier spectrum computation of a scalar field.
+    """
+
+    def __init__(self, field, prefix=None, **kwds):
+        """
+        Constructor for Spectrum operator
+        @param[in] field : field to compute
+        """
+        super(Spectrum, self).__init__(variables=[field], **kwds)
+        self.field = field
+        self.input = [field]
+        self._prefix = prefix
+
+    def discretize(self):
+        super(Spectrum, self)._fftw_discretize()
+
+    @debug
+    @opsetup
+    def setup(self, rwork=None, iwork=None):
+        self.discrete_op = FFTSpectrum(self.discreteFields[self.field],
+                                       method=self.method,
+                                       prefix=self._prefix)
+        self._is_uptodate = True

HySoP/hysop/operator/tests/test_spectrum.py

+from hysop.domain.box import Box
+from hysop.constants import np, HDF5
+from hysop.fields.continuous import Field
+from hysop.operator.spectrum import Spectrum
+from hysop.tools.parameters import Discretization, IOParams
+from hysop.problem.simulation import Simulation
+from hysop.operator.hdf_io import HDF_Writer
+pi = np.pi
+sin = np.sin
+cos = np.cos
+
+
+def computeScal(res, x, y, z, t):
+    res[0][...] = z * sin((3*pi*x)*(2*pi*y))
+    return res
+
+
+def test_spectrum():
+    dom = Box(origin=[0., 0., 0.], length=[1., 1., 1.])
+    field = Field(domain=dom, name='Field',
+                  is_vector=False, formula=computeScal)
+    d3D = Discretization([257, 257, 257])
+
+    op = Spectrum(field, discretization=d3D)
+    op.discretize()
+    op.setup()
+    topo = op.discreteFields[field].topology
+    field.initialize(topo=topo)
+
+    prt = HDF_Writer(variables={field: topo, },
+                     io_params=IOParams(frequency=1,
+                                        filename='field',
+                                        fileformat=HDF5))
+    prt.discretize()
+    prt.setup()
+
+    simu = Simulation(nbIter=1)
+    simu.initialize()
+    op.apply(simu)
+    prt.apply(simu)
+
+
+if __name__ == '__main__':
+    test_spectrum()

HySoP/src/fftw/fft3d.f90

@@ -26,7 +26,7 @@ module fft3d
        getParamatersTopologyFFTW3d,filter_poisson_3d,filter_curl_diffusion_3d, &
        init_r2c_3d_many, r2c_3d_many, c2r_3d_many, filter_diffusion_3d_many,&
        filter_poisson_3d_many, filter_diffusion_3d, filter_curl_3d, filter_projection_om_3d,&
-       filter_multires_om_3d, filter_pressure_3d
+       filter_multires_om_3d, filter_pressure_3d, r2c_3d_scal, filter_spectrum_3d
 
   !> plan for fftw "c2c" forward or r2c transform
   type(C_PTR) :: plan_forward1, plan_forward2, plan_forward3
@@ -1038,4 +1038,68 @@ contains
   end subroutine getParamatersTopologyFFTW3d
 
 
+  !> forward transform - The result is saved in local buffers
+  !!  @param[in] field 3d scalar field, x-component of the input vector field
+  !!  @param[in] ghosts, number of points in the ghost layer of input fields.
+  subroutine r2c_3d_scal(field, ghosts)
+
+    real(mk),dimension(:,:,:),intent(in) :: field
+    integer, dimension(3), intent(in) :: ghosts
+    !real(8) :: start
+    integer(C_INTPTR_T) :: i,j,k, ig, jg, kg
+
+    ! ig, jg, kg are used to take into account
+    ! ghost points in input fields
+    ! init
+    do k =1, local_resolution(c_Z)
+       kg = k + ghosts(c_Z)
+       do j = 1, fft_resolution(c_Y)
+          jg = j + ghosts(c_Y)
+          do i = 1, fft_resolution(c_X)
+             ig = i + ghosts(c_X)
+             rdatain1(i,j,k) = field(ig,jg,kg)
+          end do
+       end do
+    end do
+
+    ! compute transforms for each component
+    !start = MPI_WTIME()
+    call fftw_mpi_execute_dft_r2c(plan_forward1, rdatain1, dataout1)
+    !!print *, "r2c time = ", MPI_WTIME() - start
+
+  end subroutine r2c_3d_scal
+
+  !> Compute spectrum of the given data
+  subroutine filter_spectrum_3d(spectrum,wavelengths,length)
+
+    real(mk),dimension(:),intent(out) :: spectrum
+    real(mk),dimension(:),intent(out) :: wavelengths
+    real(mk),intent(in) :: length
+    integer(C_INTPTR_T) :: i,j,k
+    real(mk) :: c
+    real(mk) :: kk,dk,kc,eps
+    integer(C_INTPTR_T) :: ik
+    spectrum = 0
+    dk = 2.0_mk*pi/length
+    kc = pi*fft_resolution(c_X)/length
+    eps=kc/1000000.0_mk
+
+    !! mind the transpose -> index inversion between y and z
+    c = 1._mk/real(fft_resolution(c_Z)*fft_resolution(c_Y)*fft_resolution(c_X),mk)
+    c = c * c
+    do j = 1,local_resolution(c_Y)
+       do k = 1, fft_resolution(c_Z)
+          do i = 1, local_resolution(c_X)
+             kk=sqrt(kx(i)**2+ky(j)**2+kz(k)**2)
+             if ((kk.gt.eps).and.(kk.le.kc)) then
+                ik=1+int(kk/dk+0.5_mk)
+                spectrum(ik) = spectrum(ik) + real(dataout1(i,j,k)*conjg(dataout1(i,j,k)), mk) * c
+             end if
+          end do
+       end do
+    end do
+    wavelengths(:) = kx
+
+  end subroutine filter_spectrum_3d
+
 end module fft3d