(expansion to a larger supercell).master
parent
8c37b36673
commit
733c518735
1 changed files with 179 additions and 0 deletions
@ -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) |
||||
|
Loading…
Reference in new issue