From 733c518735c64a61c5ac6867e192388d7226bf4e Mon Sep 17 00:00:00 2001 From: wirawan Date: Wed, 24 Feb 2010 14:27:23 +0000 Subject: [PATCH] * Added: density refitter onto a different FFT grid and/or supercell size (expansion to a larger supercell). --- math/fft.py | 179 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 179 insertions(+) create mode 100644 math/fft.py diff --git a/math/fft.py b/math/fft.py new file mode 100644 index 0000000..d07e5cb --- /dev/null +++ b/math/fft.py @@ -0,0 +1,179 @@ +# $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) +