37 lines
1.4 KiB
Python
37 lines
1.4 KiB
Python
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)
|
|
|
|
|