New fftw module

python/cosmotool/__init__.py

@@ -1,7 +1,12 @@
 from _cosmotool import *
 from _project import *
-from grafic import writeGrafic, writeWhitePhase, readGrafic, readWhitePhase
-from borg import read_borg_vol
-from cic import cicParticles
-from simu import loadRamsesAll, simpleWriteGadget, SimulationBare
-from timing import timeit, timeit_quiet
+from .grafic import writeGrafic, writeWhitePhase, readGrafic, readWhitePhase
+from .borg import read_borg_vol
+from .cic import cicParticles
+from .simu import loadRamsesAll, simpleWriteGadget, SimulationBare
+from .timing import timeit, timeit_quiet
+
+try:
+  from .fftw import CubeFT
+except ImportError:
+  print("No FFTW support")

python/cosmotool/fftw.py

@@ -0,0 +1,36 @@
+import pyfftw
+import multiprocessing
+import numpy as np
+import numexpr as ne
+
+class CubeFT(object):
+  def __init__(self, L, N, max_cpu=-1):
+  
+      self.N = N
+      self.align = pyfftw.simd_alignment
+      self.L = L
+      self.max_cpu = multiprocessing.cpu_count() if max_cpu < 0 else max_cpu
+      self._dhat = pyfftw.n_byte_align_empty((self.N,self.N,self.N/2+1), self.align, dtype='complex64')
+      self._density = pyfftw.n_byte_align_empty((self.N,self.N,self.N), self.align, dtype='float32')
+      self._irfft = pyfftw.FFTW(self._dhat, self._density, axes=(0,1,2), direction='FFTW_BACKWARD', threads=self.max_cpu, normalize_idft=False)
+      self._rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False)
+
+  def rfft(self):
+      return ne.evaluate('c*a', out=self._dhat, local_dict={'c':self._rfft(normalise_idft=False),'a':(self.L/self.N)**3}, casting='unsafe')
+      
+  def irfft(self):
+      return ne.evaluate('c*a', out=self._density, local_dict={'c':self._irfft(normalise_idft=False),'a':(1/self.L)**3}, casting='unsafe')
+   
+  def get_dhat(self):
+      return self._dhat
+  def set_dhat(self, in_dhat):
+      self._dhat[:] = in_dhat 
+  dhat = property(get_dhat, set_dhat, None)
+      
+  def get_density(self):
+      return self._density      
+  def set_density(self, d):
+      self._density[:] = d 
+  density = property(get_density, set_density, None)
+      
+