wpylib/math/fft.py

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