# \$Id: fft.py,v 1.2 2011-10-06 19:14:48 wirawan Exp \$ # # wpylib.math.fft module # Created: 20100205 # Wirawan Purwanto # """ wpylib.math.fft FFT support. """ import sys import numpy import numpy.fft from wpylib.text_tools import slice_str from wpylib.generators import all_combinations # The minimum and maximum grid coordinates for a given FFT grid size (Gsize). # In multidimensional FFT grid, Gsize should be a numpy array. fft_grid_bounds = lambda Gsize : ( -(Gsize // 2), -(Gsize // 2) + Gsize - 1 ) """ Notes on FFT grid ranges: The fft_grid_ranges* functions define the negative and positive frequency domains on the FFT grid. Unfortunately we cannot copy an FFT grid onto another with a different grid size in single statement like: out_grid[gmin:gmax:gstep] = in_grid[:] The reason is: because gmin < gmax, python does not support such a wrapped-around array slice. The slice [gmin:gmax:gstep] will certainly result in an empty slice. To do this, we define two functions below. First, fft_grid_ranges1 generates the ranges for each dimension, then fft_grid_ranges itself generates all the combination of ranges (which cover all combinations of positive and negative frequency domains for all dimensions.) For a (5x8) FFT grid, we will have Gmin = (-2, -4) Gmax = (2, 3) Gstep = (1,1) for simplicity In this case, fft_grid_ranges1(Gmin, Gmax, Gstep) will yield [ (-2::1, 0:3:1), # negative and frequency ranges for x dimension (-4::1, 0:4:1) # negative and frequency ranges for y dimension ] [Here a:b:c is the slice(a,b,c) object in python.] All the quadrant combinations will be generated by fft_grid_ranges, which in this case is: [ (-2::1, -4::1), # -x, -y (0:3:1, -4::1), # +x, -y (-2::1, 0:4:1), # -x, +y (0:3:1, 0:4:1), # +x, +y ] """ fft_grid_ranges1 = lambda Gmin, Gmax, Gstep : \ [ (slice(gmin, None, gstep), slice(0, gmax+1, gstep)) for (gmin, gmax, gstep) in zip(Gmin, Gmax, Gstep) ] fft_grid_ranges = lambda Gmin, Gmax, Gstep : \ all_combinations(fft_grid_ranges1(Gmin, Gmax, Gstep)) class fft_grid(object): """A class describing a N-dimensional grid for plane wave (or real-space) basis. In this version, the grid is centered at (0,0,...) coordinate. To actually create a grid, use the new_dens() method. """ dtype = complex def __init__(self, Gsize=None, Gmin=None, Gmax=None, dtype=None): """Creates a new grid descriptor. There are two possible methods, and you must choose either one for initialization: * Gsize = an N-dimensional array (list, tuple, ndarray) specifying the number of grid points in each dimension. or * Gmin, Gmax = a pair of N-dimensional arrays (list, tuple, ndarray) specifying the smallest (most negative) and largest (most positive) coordinates in each dimension. The grid size will be specified to fit this range. """ from numpy import maximum if Gsize != None: self.Gsize = numpy.array(Gsize, dtype=int) (self.Gmin, self.Gmax) = fft_grid_bounds(self.Gsize) elif Gmin != None and Gmax != None: self.Gmin = numpy.array(Gmin, dtype=int) self.Gmax = numpy.array(Gmax, dtype=int) # Figure out the minimum grid size to fit this data: Gsize_min = abs(self.Gmin) * 2 Gsize_max = abs(self.Gmax) * 2 + (abs(self.Gmax) % 2) Gsize_def = self.Gmax - self.Gmin + 1 self.Gsize = maximum(maximum(Gsize_min, Gsize_max), Gsize_def) else: raise ValueError, \ "Either Gsize or (Gmin,Gmax) parameters have to be specified." if dtype != None: self.dtype = dtype self.ndim = len(self.Gsize) def new_dens(self, zero=False, dtype=None): """Creates a new N-dimensional array (grid).""" if dtype == None: dtype = self.dtype if zero: return numpy.zeros(self.Gsize, dtype=dtype) else: return numpy.empty(self.Gsize, dtype=dtype) def check_index(self, G): """Check if an index is valid according to Gmin, Gmax boundary.""" return numpy.all(self.Gmin <= G) and numpy.all(G <= self.Gmax) def fft_r2g(dens): """Do real-to-G space transformation. According to our covention, this transformation gets the 1/Vol prefactor.""" dens_G = numpy.fft.fftn(dens) dens_G *= (1.0 / numpy.prod(dens.shape)) return dens_G def fft_g2r(dens): """Do G-to-real space transformation. According to our covention, this transformation does NOT get the 1/Vol prefactor.""" dens_G = numpy.fft.ifftn(dens) dens_G *= numpy.prod(dens.shape) return dens_G def refit_grid(dens, gridsize, supercell=None, debug=0, debug_grid=False): """Refit a given density (field) to a new grid size (`gridsize'), optionally replicating in each direction by `supercell'. This function is useful for refitting/interpolation (by specifying a larger grid), low-pass filter (by specifying a smaller grid), and/or replicating a given data to construct a supercell. The dens argument is the original data on a `ndim'-dimensional FFT grid. The gridsize is an ndim-integer tuple defining the size of the new FFT grid. The supercell is an ndim-integer tuple defining the multiplicity of the new data in each direction; default: (1, 1, ...). """ from numpy import array, ones, zeros from numpy import product, minimum #from numpy.fft import fftn, ifftn # Input grid LL = array(dens.shape) ndim = len(LL) if supercell == None: supercell = ones(1, dtype=int) elif ndim != len(supercell): raise ValueError, "Incorrect supercell dimension" if ndim != len(gridsize): raise ValueError, "Incorrect gridsize dimension" #Lmin = -(LL // 2) #Lmax = Lmin + LL - 1 #Lstep = ones(LL.shape, dtype=int) # Output grid supercell = array(supercell) KK = array(gridsize) # Input grid specification for copying amplitudes: # Only this big of the subgrid from the original data will be copied: IG_size = minimum(KK // supercell, LL) (IG_min, IG_max) = fft_grid_bounds(IG_size) IG_step = ones(IG_size.shape, dtype=int) IG_ranges = fft_grid_ranges(IG_min, IG_max, IG_step) # FIXME: must check where the boundary of the nonzero G components and # warn user if we remove high frequency components # Output grid specification for copying amplitudes: # - grid stepping is identical to supercell multiplicity in each dimension # - the bounds must be commensurate to supercell steps and must the # steps must pass through (0,0,0) OG_min = IG_min * supercell OG_max = IG_max * supercell OG_step = supercell OG_ranges = fft_grid_ranges(OG_min, OG_max, OG_step) # Now form the density in G space, and copy the amplitudes to the new # grid (still in G space) if debug_grid: global dens_G global newdens_G dens_G = fft_r2g(dens) newdens_G = zeros(gridsize, dtype=dens_G.dtype) for (in_range, out_range) in zip(IG_ranges, OG_ranges): # Copies the data to the new grid, in `quadrant-by-quadrant' manner: if debug >= 1: print "G[%s] = oldG[%s]" % (slice_str(out_range), slice_str(in_range)) if debug >= 10: print dens_G[in_range] newdens_G[out_range] = dens_G[in_range] # Special case: if input size is even and the output grid is larger, # we will have to split the center bin (i.e. the highest frequency) # because it stands for both the exp(-i phi_max) and exp(+i phi_max) # (Nyquist) terms. # See: http://www.elisanet.fi/~d635415/webroot/MatlabOctaveBlocks/mn_FFT_interpolation.m select_slice = lambda X, dim : \ tuple([ slice(None) ] * dim + [ X ] + [ slice(None) ] * (ndim-dim-1)) for dim in xrange(ndim): if IG_size[dim] % 2 == 0 \ and KK[dim] > IG_size[dim] * supercell[dim]: Ny_ipos = select_slice(OG_max[dim]+1, dim) Ny_ineg = select_slice(OG_min[dim], dim) if debug > 1: print "dim", dim, ": insize=", IG_size[dim], ", outsize=", KK[dim] print "ipos = ", Ny_ipos print "ineg = ", Ny_ineg if debug > 10: print "orig dens value @ +Nyquist freq:\n" print newdens_G[Ny_ipos] newdens_G[Ny_ipos] += newdens_G[Ny_ineg] * 0.5 newdens_G[Ny_ineg] *= 0.5 return fft_g2r(newdens_G)