My tools of the trade for python programming.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

180 lines
6.2 KiB

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