# $Id: fft.py,v 1.1 2010-02-24 14:27:23 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 ndgative 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)) 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)