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.

230 lines
8.1 KiB

# $Id:,v 1.2 2011-10-06 19:14:48 wirawan Exp $
# wpylib.math.fft module
# Created: 20100205
# Wirawan Purwanto
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
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
* Gsize = an N-dimensional array (list, tuple, ndarray) specifying
the number of grid points in each dimension.
* 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)
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)
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 /
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
dens_G = numpy.fft.ifftn(dens)
dens_G *=
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:
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)