* Added: density refitter onto a different FFT grid and/or supercell size

(expansion to a larger supercell).

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)
+