convolution.py | convolution.py | |||
---|---|---|---|---|
# coding=utf-8 | # coding=utf-8 | |||
import numpy as np | import numpy as np | |||
try: | try: | |||
from pyfftw.interfaces.numpy_fft import rfftn, irfftn, fftshift | from pyfftw.interfaces.numpy_fft import rfftn, irfftn, fftshift, rfft2, irfft2 | |||
except ImportError: | except ImportError: | |||
from numpy.fft import rfftn, irfftn, fftshift | from numpy.fft import rfftn, irfftn, fftshift, rfft2, irfft2 | |||
def convolve_nd_valid(cube, psf, compute_fourier=True): | def convolve_nd_valid(cube, psf, compute_fourier=True): | |||
""" | """ | |||
using fft convolution of n-dim arrays M & N | using fft convolution of n-dim arrays M & N | |||
Returns a cube of dimensions (M-N+1)³ | Returns a cube of dimensions (M-N+1)³ | |||
""" | """ | |||
if compute_fourier: | if compute_fourier: | |||
s1 = np.array(cube.shape) | s1 = np.array(cube.shape) | |||
s2 = np.array(psf.shape) | s2 = np.array(psf.shape) | |||
skipping to change at line 92 | skipping to change at line 92 | |||
using provided Point Spread Function. | using provided Point Spread Function. | |||
""" | """ | |||
# Compute needed padding | # Compute needed padding | |||
cubep, boxcube = padding(cube, axes=[1, 2]) | cubep, boxcube = padding(cube, axes=[1, 2]) | |||
size = np.array(np.shape(cubep)[slice(1, 3)]) | size = np.array(np.shape(cubep)[slice(1, 3)]) | |||
if compute_fourier: | if compute_fourier: | |||
psfp, boxpsf = padding(psf, axes=[1, 2]) | psfp, boxpsf = padding(psf, axes=[1, 2]) | |||
fftpsf = np.fft.rfftn(psfp, s=size, axes=[1, 2]) | fftpsf = rfftn(psfp, s=size, axes=[1, 2]) | |||
else: | else: | |||
fftpsf = psf | fftpsf = psf | |||
fftimg = np.fft.rfftn(cubep, s=size, axes=[1, 2]) | fftimg = rfftn(cubep, s=size, axes=[1, 2]) | |||
#Convolution | #Convolution | |||
fft = np.fft.fftshift(np.fft.irfftn(fftimg * fftpsf, s=size, axes=[1, 2 ]), axes=[1, 2]).real | fft = np.real(fftshift(irfftn(fftimg * fftpsf, s=size, axes=[1, 2]), ax es=[1, 2])) | |||
# Remove padding | # Remove padding | |||
cube_conv = fft[boxcube] | cube_conv = fft[boxcube] | |||
return cube_conv, fftpsf | return cube_conv, fftpsf | |||
def convolve_2d(image, psf, compute_fourier=True): | def convolve_2d(image, psf, compute_fourier=True): | |||
""" | """ | |||
Compute fft of image And of PSF | Compute fft of image And of PSF | |||
Accepts fftPSF already computed with do_PSF=False | Accepts fftPSF already computed with do_PSF=False | |||
skipping to change at line 123 | skipping to change at line 123 | |||
# Compute needed padding | # Compute needed padding | |||
cubep, boxcube = padding(image) | cubep, boxcube = padding(image) | |||
if compute_fourier: | if compute_fourier: | |||
psfp, boxpsf = padding(psf) | psfp, boxpsf = padding(psf) | |||
fftpsf = np.fft.rfft2(psfp) | fftpsf = np.fft.rfft2(psfp) | |||
else: | else: | |||
fftpsf = psf | fftpsf = psf | |||
fftimg = np.fft.rfft2(cubep) | fftimg = rfft2(cubep) | |||
# Convolution | # Convolution | |||
fft = np.fft.fftshift(np.fft.irfft2(fftimg * fftpsf)).real | fft = np.real(fftshift(irfft2(fftimg * fftpsf))) | |||
# Remove padding | # Remove padding | |||
cube_conv = fft[boxcube] | cube_conv = fft[boxcube] | |||
return cube_conv, fftpsf | return cube_conv, fftpsf | |||
def convolve_1d(data, psf, compute_fourier=True, axis=0): | def convolve_1d(data, psf, compute_fourier=True, axis=0): | |||
""" | """ | |||
Convolve data with PSF only along one dimension specified by axis (defa ult: 0) | Convolve data with PSF only along one dimension specified by axis (defa ult: 0) | |||
PSF can be the PSF data or its Fourier transform | PSF can be the PSF data or its Fourier transform | |||
skipping to change at line 151 | skipping to change at line 151 | |||
axis = np.array([axis]) | axis = np.array([axis]) | |||
# Compute needed padding | # Compute needed padding | |||
cubep, boxcube = padding(data, axes=axis) | cubep, boxcube = padding(data, axes=axis) | |||
# Get the size of the axis | # Get the size of the axis | |||
size = np.array(np.shape(cubep)[slice(axis, axis + 1)]) | size = np.array(np.shape(cubep)[slice(axis, axis + 1)]) | |||
if compute_fourier: | if compute_fourier: | |||
psfp, boxpsf = padding(psf, axes=axis) | psfp, boxpsf = padding(psf, axes=axis) | |||
fftpsf = np.fft.rfftn(psfp, s=size, axes=axis) | fftpsf = rfftn(psfp, s=size, axes=axis) | |||
else: | else: | |||
fftpsf = psf | fftpsf = psf | |||
fftimg = np.fft.rfftn(cubep, s=size, axes=axis) | fftimg = rfftn(cubep, s=size, axes=axis) | |||
# Convolution | # Convolution | |||
fft = np.fft.fftshift(np.fft.irfftn(fftimg * fftpsf, s=size, axes=axis) , axes=axis).real | fft = np.real(fftshift(irfftn(fftimg * fftpsf, s=size, axes=axis), axes =axis)) | |||
# Remove padding | # Remove padding | |||
cube_conv = fft[boxcube] | cube_conv = fft[boxcube] | |||
return cube_conv, fftpsf | return cube_conv, fftpsf | |||
def padding(cube, axes=None): | def padding(cube, axes=None): | |||
""" | """ | |||
Computes padding needed for a cube to make sure it has | Computes padding needed for a cube to make sure it has | |||
a power of 2 shape along dimensions of passed axes (default [0,1]) | a power of 2 shape along dimensions of passed axes (default [0,1]) | |||
End of changes. 10 change blocks. | ||||
10 lines changed or deleted | 10 lines changed or added | |||
This html diff was produced by rfcdiff 1.41. The latest version is available from http://tools.ietf.org/tools/rfcdiff/ |