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/