agpy 0.1 documentation

Source code for AG_fft_tools.convolve

import numpy as np
import types

try: 
    import fftw3
    has_fftw = True
[docs] def fftw2(array,nthreads=1): array = array.astype('complex') outarray = array.copy() fft_forward = fftw3.Plan(array,outarray, direction='forward', flags=['estimate'],nthreads=nthreads) fft_forward() return outarray
[docs] def ifftw2(array,nthreads=1): array = array.astype('complex') outarray = array.copy() fft_backward = fftw3.Plan(array,outarray, direction='backward', flags=['estimate'],nthreads=nthreads) fft_backward() return outarray
except ImportError: fft2 = np.fft.fft2 ifft2 = np.fft.ifft2 has_fftw = False # I performed some fft speed tests and found that scipy is slower than numpy # http://code.google.com/p/agpy/source/browse/trunk/tests/test_ffts.py # try: # #print "Attempting to import scipy. If you experience a bus error at this step, it is likely because of a bad scipy install" # import scipy # import scipy.fftpack # fft2 = scipy.fftpack.fft2 # ifft2 = scipy.fftpack.ifft2 # from scipy.special import j1 # except ImportError:
[docs]def convolve(img, kernel, crop=True, return_fft=False, fftshift=True, fft_pad=True, psf_pad=False, ignore_nan=False, quiet=False, ignore_zeros=True, min_wt=1e-8, force_ignore_zeros_off=False, normalize_kernel=np.sum, debug=False, nthreads=1): """ Convolve an image with a kernel. Returns something the size of an image. Assumes image & kernel are centered *NOTE* Order matters; the kernel should be second. Options: fft_pad - Default on. Zero-pad image to the nearest 2^n psf_pad - Default off. Zero-pad image to be at least the sum of the image sizes (in order to avoid edge-wrapping when smoothing) crop - Default on. Return an image of the size of the largest input image. If the images are asymmetric in opposite directions, will return the largest image in both directions. return_fft - Return the FFT instead of the convolution. Useful for making PSDs. fftshift - If return_fft on, will shift & crop image to appropriate dimensions ignore_nan - attempts to re-weight assuming NAN values are meant to be ignored, not treated as zero. ignore_zeros - Ignore the zero-pad-created zeros. Desirable if you have periodic boundaries on a non-2^n grid force_ignore_zeros_off - You can choose to turn off the ignore-zeros when padding, but this is only recommended for purely debug purposes min_wt - If ignoring nans/zeros, force all grid points with a weight less than this value to NAN (the weight of a grid point with *no* ignored neighbors is 1.0) normalize_kernel - if specified, function to divide kernel by to normalize it nthreads - if fftw3 is installed, can specify the number of threads to allow FFTs to use. Probably only helpful for large arrays """ # replace fft2 if has_fftw so that nthreads can be passed if has_fftw: def fft2(*args, **kwargs): return fftw2(*args, nthreads=nthreads, **kwargs) def ifft2(*args, **kwargs): return ifftw2(*args, nthreads=nthreads, **kwargs) # mask catching if hasattr(img,'mask'): mask = img.mask img = np.array(img) img[mask] = np.nan # NAN catching nanmaskimg = img!=img img[nanmaskimg] = 0 if nanmaskimg.sum() > 0 and not ignore_nan and not quiet: print "Warning: NOT ignoring nan values even though they are present (they are treated as 0)" if (psf_pad or fft_pad) and not ignore_zeros and not force_ignore_zeros_off and not quiet: print "Warning: when psf_pad or fft_pad are enabled, ignore_zeros is forced on" ignore_zeros=True elif force_ignore_zeros_off: ignore_zeros=False if normalize_kernel: # try this. If a function is not passed, the code will just crash... I think type checking would be better but PEPs say otherwise... kernel /= normalize_kernel(kernel) if debug: print "Status: ignore_zeros=",ignore_zeros," force_ignore_zeros_off=",force_ignore_zeros_off," psf_pad=",psf_pad," fft_pad=",fft_pad," normalize_kernel=",normalize_kernel imgshape = img.shape kernshape = kernel.shape # find ideal size (power of 2) for fft. Can add shapes because they are tuples if fft_pad: if psf_pad: # add the X dimensions and Y dimensions and then take the max (bigger) fsize = 2**np.ceil(np.log2(np.max(np.array(imgshape)+np.array(kernshape)))) else: # add the shape lists (max of a list of length 4) (smaller) fsize = 2**np.ceil(np.log2(np.max(imgshape+kernshape))) newshape = np.array([fsize,fsize]) else: if psf_pad: newshape = np.array(imgshape)+np.array(kernshape) # just add the biggest dimensions else: newshape = np.array([np.max([imgshape[0],kernshape[0]]),np.max([imgshape[1],kernshape[1]])]) centerx, centery = newshape/2. imgquarter1x, imgquarter1y = centerx - imgshape[0]/2.,centery - imgshape[1]/2. imgquarter3x, imgquarter3y = centerx + imgshape[0]/2.,centery + imgshape[1]/2. kernelquarter1x, kernelquarter1y = centerx - kernshape[0]/2.,centery - kernshape[1]/2. kernelquarter3x, kernelquarter3y = centerx + kernshape[0]/2.,centery + kernshape[1]/2. bigimg = np.zeros(newshape,dtype=np.complex128) bigkernel = np.zeros(newshape,dtype=np.complex128) bigimg[imgquarter1x:imgquarter3x,imgquarter1y:imgquarter3y] = img bigkernel[kernelquarter1x:kernelquarter3x,kernelquarter1y:kernelquarter3y] = kernel imgfft = fft2(bigimg) kernfft = fft2(bigkernel) fftmult = imgfft*kernfft if debug: print "Kernel sum: %g Bigkernel sum: %g" % (kernel.sum(), bigkernel.sum()) if ignore_nan or ignore_zeros: bigimwt = np.ones(newshape,dtype=np.complex128) if ignore_zeros: bigimwt[:] = 0.0 bigimwt[imgquarter1x:imgquarter3x,imgquarter1y:imgquarter3y] = 1.0-nanmaskimg*ignore_nan wtfft = fft2(bigimwt) wtfftmult = wtfft*kernfft/kernel.sum() wtfftsm = ifft2(wtfftmult) pixel_weight = np.fft.fftshift(wtfftsm).real if debug: print "Pixel weight: %g bigimwt: %g wtfft sum: %g kernfft sum: %g wtfftsm sum: %g" % (pixel_weight.sum(), bigimwt.sum(), wtfft.sum(), kernfft.sum(), wtfftsm.sum()) if debug: print "img shape:",imgshape," kern shape:",kernshape," newshape:",newshape,\ " imgquarter1x,3x,1y,3y:",imgquarter1x,imgquarter3x,imgquarter1y,imgquarter3y, \ " kernelquarter1x,3x,1y,3y:",kernelquarter1x,kernelquarter3x,kernelquarter1y,kernelquarter3y if return_fft: if fftshift: # default on if crop: return np.fft.fftshift(fftmult)[ imgquarter1x:imgquarter3x, imgquarter1y:imgquarter3y ] else: return np.fft.fftshift(fftmult) else: return fftmult if ignore_nan or ignore_zeros: rifft = np.fft.fftshift( ifft2( fftmult ) ) / pixel_weight rifft[pixel_weight < min_wt] = np.nan else: rifft = np.fft.fftshift( ifft2( fftmult ) ) if crop: result = rifft[ imgquarter1x:imgquarter3x, imgquarter1y:imgquarter3y ].real return result else: return rifft.real
[docs]def smooth(image, kernelwidth=3, kerneltype='gaussian', trapslope=None, silent=True, psf_pad=True, interp_nan=False, nwidths='max', min_nwidths=6, return_kernel=False, normalize_kernel=np.sum, **kwargs): """ Returns a smoothed image using a gaussian, boxcar, or tophat kernel Options: kernelwidth - width of kernel in pixels (see definitions below) kerneltype - gaussian, boxcar, or tophat. For a gaussian, uses a gaussian with sigma = kernelwidth (in pixels) out to [nwidths]-sigma A boxcar is a kernelwidth x kernelwidth square A tophat is a flat circle with radius = kernelwidth Default options: psf_pad = True - will pad the input image to be the image size + PSF. Slows things down but removes edge-wrapping effects (see convolve) This option should be set to false if the edges of your image are symmetric. interp_nan = False - Will replace NaN points in an image with the smoothed average of its neighbors (you can still simply ignore NaN values by setting ignore_nan=True but leaving interp_nan=False) silent = True - turn it off to get verbose statements about kernel types return_kernel = False - If set to true, will return the kernel as the second return value nwidths = 'max' - number of kernel widths wide to make the kernel. Set to 'max' to match the image shape, otherwise use any integer min_nwidths = 6 - minimum number of gaussian widths to make the kernel (the kernel will be larger than the image if the image size is < min_widths*kernelsize) normalize_kernel - Should the kernel preserve the map sum (i.e. kernel.sum() = 1) or the kernel peak (i.e. kernel.max() = 1) ? Must be a *function* that can operate on a numpy array Note that the kernel is forced to be even sized on each axis to assure no offset when smoothing. """ if (kernelwidth*min_nwidths > image.shape[0] or kernelwidth*min_nwidths > image.shape[1]): nwidths = min_nwidths if (nwidths!='max'):# and kernelwidth*nwidths < image.shape[0] and kernelwidth*nwidths < image.shape[1]): dimsize = np.ceil(kernelwidth*nwidths) dimsize += dimsize % 2 yy,xx = np.indices([dimsize,dimsize]) szY,szX = dimsize,dimsize else: szY,szX = image.shape szY += szY % 2 szX += szX % 2 yy,xx = np.indices([szY,szX]) shape = (szY,szX) if not silent: print "Kernel size set to ",shape rr = np.sqrt((xx-szX/2.)**2+(yy-szY/2.)**2) if kerneltype == 'gaussian': kernel = np.exp(-(rr**2)/(2.*kernelwidth**2)) kernel /= normalize_kernel(kernel) #/ (kernelwidth**2 * (2*np.pi)) # if kernelwidth != np.round(kernelwidth): # print "Rounding kernel width to %i pixels" % np.round(kernelwidth) # kernelwidth = np.round(kernelwidth) elif kerneltype == 'boxcar': if not silent: print "Using boxcar kernel size %i" % np.ceil(kernelwidth) kernel = np.ones([np.ceil(kernelwidth),np.ceil(kernelwidth)],dtype='float64') / kernelwidth**2 elif kerneltype == 'tophat': kernel = np.zeros(shape,dtype='float64') kernel[rr<kernelwidth] = 1.0 # normalize kernel /= normalize_kernel(kernel) elif kerneltype == 'brickwall': if not silent: print "Smoothing with a %i pixel airy function" % kernelwidth # airy function is first bessel(x) / x [like the sinc] print "WARNING: I think the Airy should be (2*Besel(x)/x)^2?" # http://en.wikipedia.org/wiki/Airy_disk kernel = j1(rr/kernelwidth) / (rr/kernelwidth) # fix NAN @ center kernel[rr==0] = 0.5 # normalize - technically, this should work, but practically, flux is GAINED in some images. kernel /= normalize_kernel(kernel) elif kerneltype == 'trapezoid': if trapslope: zz = rr.max()-(rr*trapslope) zz[zz<0] = 0 zz[rr<kernelwidth] = 1.0 kernel = zz/zz.sum() else: if not silent: print "trapezoid function requires a slope" if not silent: print "Kernel of type %s normalized with %s has peak %g" % (kerneltype, normalize_kernel, kernel.max()) bad = (image != image) temp = image.copy() # to preserve NaN values # convolve does this already temp[bad] = 0 # kwargs parsing to avoid duplicate keyword passing if not kwargs.has_key('ignore_zeros'): kwargs['ignore_zeros']=True if not kwargs.has_key('ignore_nan'): kwargs['ignore_nan']=interp_nan # No need to normalize - normalization is dealt with in this code temp = convolve(temp,kernel,psf_pad=psf_pad, normalize_kernel=False, **kwargs) if interp_nan is False: temp[bad] = image[bad] if temp.shape != image.shape: raise ValueError("Output image changed size; this is completely impossible.") if return_kernel: return temp,kernel else: return temp