wpylib/math/fft.py

# $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)