* Added: density refitter onto a different FFT grid and/or supercell size
(expansion to a larger supercell).
This commit is contained in:
wirawan
179
math/fft.py
Normal file
179
math/fft.py
Normal file
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user