def hg_conv(f, g): ub = tuple(array(f.shape) + array(g.shape) - 1); fh = rfftn(cpad(f,ub)); gh = rfftn(cpad(g,ub)); res = ifftshift(irfftn(fh * sp.conjugate(gh))); del fh, gh; return res;
def computeDisplacement(arry1, arry2):
"""
Compute an optimal displacement between two ndarrays.
Finds the displacement between two ndimensional arrays. Arrays must be
of the same size. Algorithm uses a cross correlation, computed efficiently
through an n-dimensional fft.
Parameters
----------
arry1 : ndarray
The first array
arry2 : ndarray
The second array
"""
from numpy.fft import rfftn, irfftn
from numpy import unravel_index, argmax
# compute real-valued cross-correlation in fourier domain
s = arry1.shape
f = rfftn(arry1)
f *= rfftn(arry2).conjugate()
c = abs(irfftn(f, s))
# find location of maximum
inds = unravel_index(argmax(c), s)
# fix displacements that are greater than half the total size
pairs = zip(inds, arry1.shape)
# cast to basic python int for serialization
adjusted = [int(d - n) if d > n // 2 else int(d) for (d, n) in pairs]
return adjusted
def customfftconvolve(in1, in2, mode="full", types=('','')):
""" Pretty much the same as original fftconvolve, but supports
having operands as fft already
"""
in1 = asarray(in1)
in2 = asarray(in2)
if in1.ndim == in2.ndim == 0: # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("in1 and in2 should have the same dimensionality")
elif in1.size == 0 or in2.size == 0: # empty arrays
return array([])
s1 = array(in1.shape)
s2 = array(in2.shape)
complex_result = False
#complex_result = (np.issubdtype(in1.dtype, np.complex) or
# np.issubdtype(in2.dtype, np.complex))
shape = s1 + s2 - 1
if mode == "valid":
_check_valid_mode_shapes(s1, s2)
# Speed up FFT by padding to optimal size for FFTPACK
fshape = [_next_regular(int(d)) for d in shape]
fslice = tuple([slice(0, int(sz)) for sz in shape])
if not complex_result:
if types[0] == 'fft':
fin1 = in1#_unfold_fft(in1, fshape)
else:
fin1 = rfftn(in1, fshape)
if types[1] == 'fft':
fin2 = in2#_unfold_fft(in2, fshape)
else:
fin2 = rfftn(in2, fshape)
ret = irfftn(fin1 * fin2, fshape)[fslice].copy()
else:
if types[0] == 'fft':
fin1 = _unfold_fft(in1, fshape)
else:
fin1 = fftn(in1, fshape)
if types[1] == 'fft':
fin2 = _unfold_fft(in2, fshape)
else:
fin2 = fftn(in2, fshape)
ret = ifftn(fin1 * fin2)[fslice].copy()
if mode == "full":
return ret
elif mode == "same":
return _centered(ret, s1)
elif mode == "valid":
return _centered(ret, s1 - s2 + 1)
else:
raise ValueError("Acceptable mode flags are 'valid',"
" 'same', or 'full'.")
def _cpu_init(self):
self.cpu_data = {}
c = self.cpu_data
d = self.data
c['rcore'] = d['rcore'].array
c['rsurf'] = d['rsurf'].array
c['im_lsurf'] = d['lsurf'].array
c['lsurf'] = np.zeros_like(c['rcore'])
c['clashvol'] = np.zeros_like(c['rcore'])
c['intervol'] = np.zeros_like(c['rcore'])
c['interspace'] = np.zeros_like(c['rcore'], dtype=np.int64)
# complex arrays
c['ft_shape'] = list(d['shape'])
c['ft_shape'][-1] = d['shape'][-1]//2 + 1
c['ft_lsurf'] = np.zeros(c['ft_shape'], dtype=np.complex128)
c['ft_rcore'] = np.zeros(c['ft_shape'], dtype=np.complex128)
c['ft_rsurf'] = np.zeros(c['ft_shape'], dtype=np.complex128)
# initial calculations
c['ft_rcore'] = rfftn(c['rcore'])
c['ft_rsurf'] = rfftn(c['rsurf'])
c['rotmat'] = np.asarray(self.rotations, dtype=np.float64)
c['weights'] = np.asarray(self.weights, dtype=np.float64)
c['nrot'] = d['nrot']
c['shape'] = d['shape']
c['max_clash'] = d['max_clash']
c['min_interaction'] = d['min_interaction']
c['vlength'] = int(np.linalg.norm(self.ligand.coor - \
self.ligand.center, axis=1).max() + \
self.interaction_radius + 1.5)/self.voxelspacing
c['origin'] = d['origin']
# SAXS arrays
c['q'] = d['q']
c['targetIq'] = d['targetIq']
c['sq'] = d['sq']
c['base_Iq'] = d['base_Iq']
c['fifj'] = d['fifj']
c['rind'] = d['rind']
c['lind'] = d['lind']
c['rxyz'] = d['rxyz']
c['lxyz'] = d['lxyz']
c['chi2'] = d['chi2']
c['best_chi2'] = d['best_chi2']
c['rot_ind'] = np.zeros(d['shape'], dtype=np.int32)
c['Iq'] = np.zeros_like(c['targetIq'])
c['tmplxyz'] = np.zeros_like(c['lxyz'])
def SpatialCorrelationFunctionA(Field1, Field2):
"""
Designed for Periodic Boundary Condition.
Corr_12(r) = <Phi_1(r)Phi_2(0)>
Corr_12(k) = Phi_1(k)* Phi_2(k)/V
"""
V = float(numpy.array(Field1.shape).prod())
KField1 = fft.rfftn(Field1).conj()
KField2 = fft.rfftn(Field2)
KCorr = KField1 * KField2 / V
Corr = fft.irfftn(KCorr)
return Corr
def fftconvolve(in1, in2):
"""Convolve two N-dimensional arrays using FFT.
This is a modified version of the scipy.signal.fftconvolve.
The new feature is derived from the fftconvolve algorithm used in the IDL package.
Parameters
----------
in1 : array_like
First input.
in2 : array_like
Second input. Should have the same number of dimensions as `in1`;
if sizes of `in1` and `in2` are not equal then `in1` has to be the
larger array.
Returns
-------
out : array
An N-dimensional array containing a subset of the discrete linear
convolution of `in1` with `in2`.
"""
in1 = asarray(in1)
in2 = asarray(in2)
if matrix_rank(in1) == matrix_rank(in2) == 0: # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("in1 and in2 should have the same rank")
elif in1.size == 0 or in2.size == 0: # empty arrays
return array([])
s1 = np.array(in1.shape)
s2 = np.array(in2.shape)
complex_result = (np.issubdtype(in1.dtype, np.complex) or
np.issubdtype(in2.dtype, np.complex))
fsize = s1
fslice = tuple([slice(0, int(sz)) for sz in fsize])
if not complex_result:
ret = irfftn(rfftn(in1, fsize) *
rfftn(in2, fsize), fsize)[fslice].copy()
ret = ret.real
else:
ret = ifftn(fftn(in1, fsize) * fftn(in2, fsize))[fslice].copy()
shift = array([int(floor(fsize[0]/2.0)), int(floor(fsize[1]/2.0))])
ret = roll(roll(ret, -shift[0], axis=0), -shift[1], axis=1)
return ret
def getCovMatrix(IQdata, lags=100, hp=False): # 0: <I1I2> # 1: <Q1Q2> # 2: <I1Q2> # 3: <Q1I2> # 4: <Squeezing> Magnitude # 5: <Squeezing> Phase lags = int(lags) I1 = np.asarray(IQdata[0]) Q1 = np.asarray(IQdata[1]) I2 = np.asarray(IQdata[2]) Q2 = np.asarray(IQdata[3]) CovMat = np.zeros([10, lags * 2 + 1]) start = len(I1) - lags - 1 stop = len(I1) + lags sI1 = np.array(I1.shape) shape0 = sI1 * 2 - 1 fshape = [_next_regular(int(d)) for d in shape0] # padding to optimal size for FFTPACK fslice = tuple([slice(0, int(sz)) for sz in shape0]) # Do FFTs and get Cov Matrix fftI1 = rfftn(I1, fshape) fftQ1 = rfftn(Q1, fshape) fftI2 = rfftn(I2, fshape) fftQ2 = rfftn(Q2, fshape) rfftI1 = rfftn(I1[::-1], fshape) # there should be a simple relationship to fftI1 rfftQ1 = rfftn(Q1[::-1], fshape) rfftI2 = rfftn(I2[::-1], fshape) rfftQ2 = rfftn(Q2[::-1], fshape) CovMat[0, :] = irfftn((fftI1 * rfftI2))[fslice].copy()[start:stop] / fshape CovMat[1, :] = irfftn((fftQ1 * rfftQ2))[fslice].copy()[start:stop] / fshape CovMat[2, :] = irfftn((fftI1 * rfftQ2))[fslice].copy()[start:stop] / fshape CovMat[3, :] = irfftn((fftQ1 * rfftI2))[fslice].copy()[start:stop] / fshape psi = (1j * (CovMat[2, :] + CovMat[3, :]) + (CovMat[0, :] - CovMat[1, :])) CovMat[4, :] = abs(psi) CovMat[5, :] = np.angle(psi) CovMat[6, :] = irfftn((fftI1 * rfftI1))[fslice].copy()[start:stop] / fshape CovMat[7, :] = irfftn((fftQ1 * rfftQ1))[fslice].copy()[start:stop] / fshape CovMat[8, :] = irfftn((fftI2 * rfftI2))[fslice].copy()[start:stop] / fshape CovMat[9, :] = irfftn((fftQ2 * rfftQ2))[fslice].copy()[start:stop] / fshape return CovMat
def _cpu_init(self):
"""Initialize all the arrays and data required for a CPU search"""
self.cpu_data = {}
c = self.cpu_data
d = self.data
# create views of data in cpu_data
c['rcore'] = d['rcore'].array
c['rsurf'] = d['rsurf'].array
c['im_lsurf'] = d['lsurf'].array
c['restraints'] = d['restraints']
# allocate arrays used for search
# real arrays
c['lsurf'] = np.zeros_like(c['rcore'])
c['clashvol'] = np.zeros_like(c['rcore'])
c['intervol'] = np.zeros_like(c['rcore'])
c['interspace'] = np.zeros_like(c['rcore'], dtype=np.int32)
c['access_interspace'] = np.zeros_like(c['rcore'], dtype=np.int32)
c['restspace'] = np.zeros_like(c['rcore'], dtype=np.int32)
c['violations'] = np.zeros((d['nrestraints'], d['nrestraints']), dtype=np.float64)
# complex arrays
c['ft_shape'] = list(d['shape'])
c['ft_shape'][-1] = d['shape'][-1]//2 + 1
c['ft_lsurf'] = np.zeros(c['ft_shape'], dtype=np.complex128)
c['ft_rcore'] = np.zeros(c['ft_shape'], dtype=np.complex128)
c['ft_rsurf'] = np.zeros(c['ft_shape'], dtype=np.complex128)
c['rotmat'] = np.asarray(self.rotations, dtype=np.float64)
c['weights'] = np.asarray(self.weights, dtype=np.float64)
c['nrot'] = d['nrot']
c['shape'] = d['shape']
c['max_clash'] = d['max_clash']
c['min_interaction'] = d['min_interaction']
# the vlenght is the longest distance from the ligand center to another
# atom. This makes the rotation of the ligand object cheaper by only
# rotation the inner part of the array where there is density
c['vlength'] = int(np.linalg.norm(self.ligand.coor - \
self.ligand.center, axis=1).max() + \
self.interaction_radius + 1.5)/self.voxelspacing
# initial calculations. Calculate the FFT of the fixed chain objects.
# This only needs to be done once before the search.
c['ft_rcore'] = rfftn(c['rcore'])
c['ft_rsurf'] = rfftn(c['rsurf'])
def multichannel_fftconvolve(x, h, mode='valid'):
x_len = x.shape[0]
num_channels = h.shape[1]
h_len = h.shape[0]
assert x.shape[1] == num_channels
assert x_len >= h_len, \
"The signal needs to be at least as long as the filter"
assert mode == 'valid'
fshape = (int(2**math.ceil(np.log2((x_len + h_len - 1)))), num_channels)
x_fft = rfftn(x, fshape)
h_fft = rfftn(h, fshape)
ret = _transformed_fft_convolve(x_fft, h_fft)
return ret[h_len-1:x_len, num_channels-1]
def SpatialCorrelationFunction(Field1,Field2):
"""
Designed for Periodic Boundary Condition.
.. math::
C_{12}(r) = <F_1(r) F_2(0)>
C_{12}(k) = F_1(k)* F_2(k)/V
"""
V = float(numpy.array(Field1.shape).prod())
KField1 = fft.rfftn(Field1).conj()
KField2 = fft.rfftn(Field2)
KCorr = KField1*KField2/V
Corr = fft.irfftn(KCorr)
return Corr
def multichannel_overlap_add_fftconvolve(x, h, mode='valid'):
"""Given a signal x compute the convolution with h using the overlap-add algorithm.
This is an fft based convolution algorithm optimized to work on a signal x
and filter, h, where x is much longer than h.
Input:
x: float array with dimensions (N, num_channel)
h: float array with dimensions (filter_len, num_channel)
mode: only accepts valid - same profile scipy.convolve
Returns:
float array of length N-filter_len+1 (for mode = valid)
"""
assert mode == 'valid'
# pad x so that the boundaries are dealt with correctly
x_len = x.shape[0]
num_channels = h.shape[1]
h_len = h.shape[0]
assert x.shape[1] == num_channels
assert x_len >= h_len, \
"The signal needs to be at least as long as the filter"
#x = np.vstack((np.zeros((h_len, num_channels)),
# x,
# np.zeros((h_len, num_channels))))
# make sure that the desired block size is long enough to capture the motif
block_size = max(2**OVERLAP_ADD_BLOCK_POWER, h_len)
N = int(2**math.ceil(np.log2(block_size+h_len-1)))
step_size = N-h_len+1
H = rfftn(h,(N,num_channels))
n_blocks = int(math.ceil(float(len(x))/step_size))
y = np.zeros((n_blocks+1)*step_size)
for block_index in xrange(n_blocks):
start = block_index*step_size
yt = irfftn( rfftn(x[start:start+step_size,:],(N, num_channels))*H,
(N, num_channels) )
y[start:start+N] += yt[:,num_channels-1]
#y = y[h_len:2*h_len+x_len-1]
if mode == 'full':
return y
elif mode == 'valid':
return y[h_len-1:x_len]
elif mode == 'same':
raise NotImplementedError, "'same' mode is not implemented"
def cross_correlation(seqs):
# deal with the shape, and upcast to the next reasonable shape
shape = np.array(seqs.shape[1:]) + np.array(seqs.shape[1:]) - 1
fshape = [next_good_fshape(x) for x in shape]
fslice = tuple([slice(0, int(sz)) for sz in shape])
flipped_seqs_fft = np.zeros([seqs.shape[0],] + fshape[:-1] + [fshape[-1]//2+1,], dtype='complex')
for i in xrange(seqs.shape[0]):
rev_slice = tuple([i,] + [slice(None, None, -1) for sz in shape])
flipped_seqs_fft[i] = rfftn(seqs[rev_slice], fshape)
rv = np.zeros((seqs.shape[0], seqs.shape[0]), dtype='float32')
for i in xrange(seqs.shape[0]):
fft_seq = rfftn(seqs[i], fshape)
for j in xrange(i+1, seqs.shape[0]):
rv[i,j] = irfftn(fft_seq*flipped_seqs_fft[j], fshape)[fslice].max()
#print rv[i,j], correlate(seqs[i], seqs[j]).max()
return rv
def _setup_kernel(self):
if not isinstance(self.coordmap, AffineTransform):
raise ValueError, 'for FFT smoothing, need a regular (affine) coordmap'
voxels = np.indices(self.bshape).astype(np.float64)
center = np.asarray(self.bshape)/2
center = self.coordmap([center[i] for i in range(len(self.bshape))])
voxels.shape = (voxels.shape[0], np.product(voxels.shape[1:]))
X = (self.coordmap(voxels.T) - center).T
X.shape = (self.coordmap.ndim[0],) + tuple(self.bshape)
kernel = self(X)
kernel = _crop(kernel)
self.norms = {'l2':np.sqrt((kernel**2).sum()),
'l1':np.fabs(kernel).sum(),
'l1sum':kernel.sum()}
self._kernel = kernel
self.shape = (np.ceil((np.asarray(self.bshape) +
np.asarray(kernel.shape))/2)*2+2)
self.fkernel = np.zeros(self.shape)
slices = [slice(0, kernel.shape[i]) for i in range(len(kernel.shape))]
self.fkernel[slices] = kernel
self.fkernel = fft.rfftn(self.fkernel)
return kernel
def _setup_kernel(self):
if not isinstance(self.coordmap, AffineTransform):
raise ValueError('for FFT smoothing, we need a '
'regular (affine) coordmap')
# voxel indices of array implied by shape
voxels = np.indices(self.bshape).astype(np.float64)
# coordinates of physical center. XXX - why the 'floor' here?
vox_center = np.floor((np.array(self.bshape) - 1) / 2.0)
phys_center = self.coordmap(vox_center)
# reshape to (N coordinates, -1). We appear to need to assign
# to shape instead of doing a reshape, in order to avoid memory
# copies
voxels.shape = (voxels.shape[0], np.product(voxels.shape[1:]))
# physical coordinates relative to center
X = (self.coordmap(voxels.T) - phys_center).T
X.shape = (self.coordmap.ndims[0],) + tuple(self.bshape)
# compute kernel from these positions
kernel = self(X, axis=0)
kernel = _crop(kernel)
self.norms = {'l2':np.sqrt((kernel**2).sum()),
'l1':np.fabs(kernel).sum(),
'l1sum':kernel.sum()}
self._kernel = kernel
self.shape = (np.ceil((np.asarray(self.bshape) +
np.asarray(kernel.shape))/2)*2+2)
self.fkernel = np.zeros(self.shape)
slices = [slice(0, kernel.shape[i]) for i in range(len(kernel.shape))]
self.fkernel[slices] = kernel
self.fkernel = fft.rfftn(self.fkernel)
return kernel
def _setup_kernel(self):
# voxel indices of array implied by shape
voxels = np.indices(self._bshape, dtype=np.float64)
# coordinates of physical center. XXX - why the 'floor' here?
vox_center = np.floor((np.array(self._bshape) - 1) / 2.0)
phys_center = get_physical_coords(self._affine, vox_center)
# reshape to (N coordinates, -1). We appear to need to assign
# to shape instead of doing a reshape, in order to avoid memory
# copies
voxels.shape = (voxels.shape[0], np.product(voxels.shape[1:]))
# physical coordinates relative to center
X = get_physical_coords(self._affine,
voxels) - phys_center
X.shape = (self._ndims[1],) + tuple(self._bshape)
# compute kernel from these positions
kernel = self(X, axis=0)
kernel = _crop(kernel)
# compute kernel norm
self._norm = _get_kernel_norm(kernel, self._normalization)
self._kernel = kernel
self._shape = (np.ceil((np.asarray(self._bshape) +
np.asarray(kernel.shape)) / 2) * 2 + 2)
self.fkernel = np.zeros(self._shape)
slices = [slice(0, kernel.shape[i]) for i in range(kernel.ndim)]
self.fkernel[slices] = kernel
self.fkernel = npfft.rfftn(self.fkernel)
return kernel
def fftn_mpi(u, fu):
"""fft in three directions using mpi
"""
if num_processes == 1:
#fu[:] = fft(fft(rfft(u, axis=1), axis=2), axis=0)
fu[:] = rfftn(u, axes=(0,2,1))
return
# Do 2 ffts in y-z directions on owned data
#ft = fu.transpose(2,1,0)
#ft[:] = fft(rfft(u, axis=1), axis=2)
Uc_hatT[:] = rfft2(u, axes=(2,1))
## Communicating intermediate result
##rstack(ft, Uc_hatT, Np, num_processes)
#fu_send = fu.reshape((num_processes, Np, Nf, Np))
#for i in range(num_processes):
#if not i == rank:
#comm.Sendrecv_replace([fu_send[i], MPI.DOUBLE_COMPLEX], i, 0, i, 0)
#fu_send[:] = fu_send.transpose(0,3,2,1)
# Transform data to align with x-direction
for i in range(num_processes):
#U_mpi[i] = ft[:, :, i*Np:(i+1)*Np]
U_mpi[i] = Uc_hatT[:, :, i*Np:(i+1)*Np]
# Communicate all values
comm.Alltoall([U_mpi, MPI.DOUBLE_COMPLEX], [fu, MPI.DOUBLE_COMPLEX])
# Do fft for last direction
fu[:] = fft(fu, axis=0)
def get_g2(P1, P2, lags=20):
''' Returns the Top part of the G2 equation (<P1P2> - <P1><P2>)'''
lags = int(lags)
P1 = np.asarray(P1)
P2 = np.asarray(P2)
# G2 = np.zeros([lags*2-1])
start = len(P1*2-1)-lags
stop = len(P1*2-1)-1+lags
# assume I1 Q1 have the same shape
sP1 = np.array(P1.shape)
complex_result = np.issubdtype(P1.dtype, np.complex)
shape = sP1 - 1
HPfilt = (int(sP1/(lags*4))) # smallest features visible is lamda/4
# Speed up FFT by padding to optimal size for FFTPACK
fshape = [_next_regular(int(d)) for d in shape]
fslice = tuple([slice(0, int(sz)) for sz in shape])
# Pre-1.9 NumPy FFT routines are not threadsafe. For older NumPys, make
# sure we only call rfftn/irfftn from one thread at a time.
if not complex_result and _rfft_lock.acquire(False):
try:
fftP1 = rfftn(P1, fshape)
rfftP2 = rfftn(P2[::-1], fshape)
fftP1 = np.concatenate((np.zeros(HPfilt), fftP1[HPfilt:]))
rfftP2 = np.concatenate((np.zeros(HPfilt), rfftP2[HPfilt:]))
G2 = irfftn((fftP1*rfftP2))[fslice].copy()[start:stop]/len(fftP1)
return
finally:
_rfft_lock.release()
else:
# If we're here, it's either because we need a complex result, or we
# failed to acquire _rfft_lock (meaning rfftn isn't threadsafe and
# is already in use by another thread). In either case, use the
# (threadsafe but slower) SciPy complex-FFT routines instead.
# ret = ifftn(fftn(in1, fshape) * fftn(in2, fshape))[fslice].copy()
print 'Abort, reason:complex input or Multithreaded FFT not available'
if not complex_result:
pass # ret = ret.real
P12var = np.var(P1)*np.var(P2)
return G2-P12var
def GaussianRandomInitializer(gridShape, sigma=0.2, seed=None, slipSystem=None, slipPlanes=None, slipDirections=None, vacancy=None, smectic=None):
oldgrid = copy.copy(gridShape)
if len(gridShape) == 1:
gridShape = (128,)
if len(gridShape) == 2:
gridShape = (128,128)
if len(gridShape) == 3:
gridShape = (128,128,128)
""" Returns a random initial set of fields of class type PlasticityState """
if slipSystem=='gamma':
state = SlipSystemState.SlipSystemState(gridShape,slipPlanes=slipPlanes,slipDirections=slipDirections)
elif slipSystem=='betaP':
state = SlipSystemBetaPState.SlipSystemState(gridShape,slipPlanes=slipPlanes,slipDirections=slipDirections)
else:
if vacancy is not None:
state = VacancyState.VacancyState(gridShape,alpha=vacancy)
elif smectic is not None:
state = SmecticState.SmecticState(gridShape)
else:
state = PlasticityState.PlasticityState(gridShape)
field = state.GetOrderParameterField()
Ksq_prime = FourierSpaceTools.FourierSpaceTools(gridShape).kSq * (-sigma**2/4.)
if seed is None:
seed = 0
n = 0
random.seed(seed)
Ksq = FourierSpaceTools.FourierSpaceTools(gridShape).kSq.numpy_array()
for component in field.components:
temp = random.normal(scale=gridShape[0],size=gridShape)
ktemp = fft.rfftn(temp)*(sqrt(pi)*sigma)**len(gridShape)*exp(-Ksq*sigma**2/4.)
field[component] = numpy.real(fft.irfftn(ktemp))
#field[component] = GenerateGaussianRandomArray(gridShape, temp ,sigma)
n += 1
"""
t, s = LoadState("2dstate32.save", 0)
for component in field.components:
for j in range(0,32):
field[component][:,:,j] = s.betaP[component].numpy_array()
"""
## To make seed consistent across grid sizes and convergence comparison
gridShape = copy.copy(oldgrid)
if gridShape[0] != 128:
state = ResizeState(state,gridShape[0],Dim=len(gridShape))
state = ReformatState(state)
state.ktools = FourierSpaceTools.FourierSpaceTools(gridShape)
return state
def SpatialCorrelationFunctionA(Field1,Field2):
"""
Corr_12(r) = <Phi_1(r)Phi_2(0)>
Corr_12(k) = Phi_1(k)* Phi_2(k)/V
"""
dim = len(Field1.shape)
if dim == 1:
V=float(Field1.shape[0])
elif dim == 2:
V=float(Field1.shape[0]*Field1.shape[1])
elif dim == 3:
V=float(Field1.shape[0]*Field1.shape[1]*Field1.shape[2])
KField1 = fft.rfftn(Field1).conj()
KField2 = fft.rfftn(Field2)
KCorr = KField1*KField2/V
Corr = fft.irfftn(KCorr)
return Corr
def compute_fp(self):
#n_seg = matrix.shape[1] / 2**16
#idx = np.arange(2**16)
n_bands = self.processed.shape[0]
#result = np.zeros((n_seg,n_bands,61))
#take 61 instead 60 frequency bins as next step will discard last bin
fs_model_weights = self.fs_model_curve((11000./2**16) * np.arange(60))
#use rfftn
self.processed = np.c_[self.processed,np.zeros(n_bands)] #calculate fluctuation patter for each frame
self.processed = fs_model_weights * np.abs(rfftn(self.processed.reshape(n_bands,self.frame_size*2,-1).transpose(2,0,1), axes=(2,))[:,:,:60])
def GenerateGaussianRandomArray(gridShape, temp, sigma):
dimension = len(gridShape)
if dimension == 1:
kfactor = fromfunction(lambda kz: exp(-0.5*(sigma*kz)**2),[gridShape[0]/2+1,])
ktemp = fft.rfft(temp)
ktemp *= kfactor
data = fft.irfft(ktemp)
elif dimension == 2:
X,Y = gridShape
kfactor = fromfunction(lambda kx,ky: exp(-0.5*sigma**2*((kx*(kx<X/2)+(X-kx)*(kx>=X/2))**2+ky**2)),[X,Y/2+1])
ktemp = fft.rfftn(temp)
ktemp *= kfactor
data = fft.irfftn(ktemp)
elif dimension == 3:
X,Y,Z = gridShape
kfactor = fromfunction(lambda kx,ky,kz: exp(-0.5*sigma**2*( (kx*(kx<X/2)+(X-kx)*(kx>=X/2))**2 + \
(ky*(ky<Y/2)+(Y-ky)*(ky>=Y/2))**2 + kz**2)),[X,Y,Z/2+1])
ktemp = fft.rfftn(temp)
ktemp *= kfactor
data = fft.irfftn(ktemp)
return data
def zeroextract2d(N,filename):
a = fromfile(filename)
a = a.reshape(9,N,N)
b = numpy.zeros((9,N/2,N/2),float)
for i in range(9):
ka = fft.rfftn(a[i])
kb = numpy.zeros((N/2,N/4+1),complex)
kb[:N/4,:]=ka[:N/4,:N/4+1]
kb[-N/4:,:]=ka[-N/4:,:N/4+1]
b[i] = fft.irfftn(kb)
b /= 4.
b.tofile(filename.replace(str(N),str(N/2)))
def get_covariance_submatrix_full(IQdata, lags):
I1 = np.asarray(IQdata[0])
# Q1 = np.asarray(IQdata[1])
# I2 = np.asarray(IQdata[2])
Q2 = np.asarray(IQdata[3])
lags = int(lags)
start = len(I1) - lags - 1
stop = len(I1) + lags
sub_matrix = np.zeros([16, lags * 2 + 1])
sI1 = np.array(I1.shape)
shap0 = sI1*2 - 1
fshape = [_next_regular(int(d)) for d in shap0] # padding to optimal size for FFTPACK
fslice = tuple([slice(0, int(sz)) for sz in shap0]) # remove padding later
fftIQ = 4*[None]
rfftIQ = 4*[None]
for i in range(4):
fftIQ[i] = rfftn(IQdata[i], fshape)
rfftIQ[i] = rfftn(IQdata[i][::-1], fshape)
for j in range(4):
for i in range(4):
idx = i + j*4
sub_matrix[idx] = (irfftn(fftIQ[i]*rfftIQ[j]))[fslice].copy()[start:stop]/len(fftIQ[i])
return sub_matrix
def get_covariance_submatrix(IQdata, lags, q):
logging.debug('Calculating Submatrix')
I1 = np.asarray(IQdata[0])
Q1 = np.asarray(IQdata[1])
I2 = np.asarray(IQdata[2])
Q2 = np.asarray(IQdata[3])
lags = int(lags)
start = len(I1) - lags - 1
stop = len(I1) + lags
sub_matrix = np.zeros([4, lags * 2 + 1])
sI1 = np.array(I1.shape)
shape0 = sI1*2 - 1
fshape = [_next_regular(int(d)) for d in shape0] # padding to optimal size for FFTPACK
fslice = tuple([slice(0, int(sz)) for sz in shape0]) # remove padding later
fftI1 = rfftn(I1, fshape)/fshape
fftQ1 = rfftn(Q1, fshape)/fshape
rfftI2 = rfftn(I2[::-1], fshape)/fshape
rfftQ2 = rfftn(Q2[::-1], fshape)/fshape
sub_matrix[0] = irfftn((fftI1 * rfftI2))[fslice].copy()[start:stop] # <II>
sub_matrix[1] = irfftn((fftI1 * rfftQ2))[fslice].copy()[start:stop] # <IQ>
sub_matrix[2] = irfftn((fftQ1 * rfftI2))[fslice].copy()[start:stop] # <QI>
sub_matrix[3] = irfftn((fftQ1 * rfftQ2))[fslice].copy()[start:stop] # <QQ>
q.put(sub_matrix)
def dft(input, output, k, args):
# Read inputs.
if args:
data = json.loads(urllib.unquote(args))
else:
data = json.load(input)
# Make vectors.
m = np.array([vec(elem) for elem in data])
# Transform
fm = fft.rfftn(m)
# Output
json.dump([enc(v) for v in fm[:k]], output)
output.write("\n")
def getCovMatrix(I1, Q1, I2, Q2, lags=20):
# calc <I1I2>, <I1Q2>, Q1I2, Q1Q2
lags = int(lags)
I1 = np.asarray(I1)
Q1 = np.asarray(Q1)
I2 = np.asarray(I2)
Q2 = np.asarray(Q2)
CovMat = np.zeros([6, lags*2-1])
start = len(I1*2-1)-lags
stop = len(I1*2-1)-1+lags
sI1 = np.array(I1.shape)
sQ2 = np.array(Q2.shape)
shape = sI1 + sQ2 - 1
HPfilt = (int(sI1/(lags*4))) # smallest features visible is lamda/4
fshape = [_next_regular(int(d)) for d in shape] # padding to optimal size for FFTPACK
fslice = tuple([slice(0, int(sz)) for sz in shape])
# Do FFTs and get Cov Matrix
fftI1 = rfftn(I1, fshape)
fftQ1 = rfftn(Q1, fshape)
fftI2 = rfftn(I2, fshape)
fftQ2 = rfftn(Q2, fshape)
rfftI1 = rfftn(I1[::-1], fshape)
rfftQ1 = rfftn(Q1[::-1], fshape)
rfftI2 = rfftn(I2[::-1], fshape)
rfftQ2 = rfftn(Q2[::-1], fshape)
# filter frequencies outside the lags range
fftI1 = np.concatenate((np.zeros(HPfilt), fftI1[HPfilt:]))
fftQ1 = np.concatenate((np.zeros(HPfilt), fftQ1[HPfilt:]))
fftI2 = np.concatenate((np.zeros(HPfilt), fftI2[HPfilt:]))
fftQ2 = np.concatenate((np.zeros(HPfilt), fftQ2[HPfilt:]))
# filter frequencies outside the lags range
rfftI1 = np.concatenate((np.zeros(HPfilt), rfftI1[HPfilt:]))
rfftQ1 = np.concatenate((np.zeros(HPfilt), rfftQ1[HPfilt:]))
rfftI2 = np.concatenate((np.zeros(HPfilt), rfftI2[HPfilt:]))
rfftQ2 = np.concatenate((np.zeros(HPfilt), rfftQ2[HPfilt:]))
CovMat[0, :] = (irfftn((fftI1*rfftI2))[fslice].copy()[start:stop] / len(fftI1)) # 0: <I1I2>
CovMat[1, :] = (irfftn((fftQ1*rfftQ2))[fslice].copy()[start:stop] / len(fftI1)) # 1: <Q1Q2>
CovMat[2, :] = (irfftn((fftI1*rfftQ2))[fslice].copy()[start:stop] / len(fftI1)) # 2: <I1Q2>
CovMat[3, :] = (irfftn((fftQ1*rfftI2))[fslice].copy()[start:stop] / len(fftI1)) # 3: <Q1I2>
CovMat[4, :] = (abs(1j*(CovMat[2, :]+CovMat[3, :]) + (CovMat[0, :] - CovMat[1, :]))) # 4: <Squeezing> Magnitude
CovMat[5, :] = np.angle(1j*(CovMat[2, :]+CovMat[3, :]) + (CovMat[0, :] - CovMat[1, :])) # 5: <Squeezing> Angle
return CovMat
def __init__(self, psf_wrapper, flat_sky_proj):
self._psf = psf_wrapper # type: PSFWrapper
self._flat_sky_proj = flat_sky_proj
# Compute an image of the PSF on the current defined flat sky projection
interpolator = PSFInterpolator(psf_wrapper, flat_sky_proj)
psf_stamp = interpolator.point_source_image(flat_sky_proj.ra_center,
flat_sky_proj.dec_center)
# Crop the kernel at the appropriate radius for this PSF (making sure is divisible by 2)
kernel_radius_px = int(
np.ceil(self._psf.kernel_radius / flat_sky_proj.pixel_size))
pixels_to_keep = kernel_radius_px * 2
assert pixels_to_keep <= psf_stamp.shape[0] and \
pixels_to_keep <= psf_stamp.shape[1], \
"The kernel is too large with respect to the model image. Enlarge your model radius."
xoff = (psf_stamp.shape[0] - pixels_to_keep) // 2
yoff = (psf_stamp.shape[1] - pixels_to_keep) // 2
self._kernel = psf_stamp[yoff:-yoff, xoff:-xoff]
assert np.isclose(
self._kernel.sum(), 1.0, rtol=1e-2
), "Failed to generate proper kernel normalization: got _kernel.sum() = %f; expected 1.0+-0.01." % self._kernel.sum(
)
# Renormalize to exactly 1
self._kernel = self._kernel / self._kernel.sum()
self._expected_shape = (flat_sky_proj.npix_height,
flat_sky_proj.npix_width)
s1 = np.array(self._expected_shape)
s2 = np.array(self._kernel.shape)
shape = s1 + s2 - 1
self._fshape = [helper.next_fast_len(int(d)) for d in shape]
self._fslice = tuple([slice(0, int(sz)) for sz in shape])
self._psf_fft = rfftn(self._kernel, self._fshape)
def test_energy(self):
tol = 1e-10
L = 2 + rand() # domain length
a = 3 + rand() # amplitude of force
E = 4 + rand() # Young's Mod
for res in [4, 8, 16]:
area_per_pt = L / res
x = np.arange(res) * L / res
force = a * np.cos(2 * np.pi / L * x)
# theoretical FFT of force
Fforce = np.zeros_like(x)
Fforce[1] = Fforce[-1] = res / 2. * a
# theoretical FFT of disp
Fdisp = np.zeros_like(x)
Fdisp[1] = Fdisp[-1] = res / 2. * a / E * L / np.pi
# verify consistency
hs = PeriodicFFTElasticHalfSpace(res, E, L)
fforce = rfftn(force.T).T
fdisp = hs.greens_function * fforce
self.assertTrue(
Tools.mean_err(fforce, Fforce, rfft=True) < tol,
"fforce = \n{},\nFforce = \n{}".format(fforce.real, Fforce))
self.assertTrue(
Tools.mean_err(fdisp, Fdisp, rfft=True) < tol,
"fdisp = \n{},\nFdisp = \n{}".format(fdisp.real, Fdisp))
# Fourier energy
E = .5 * np.dot(Fforce / area_per_pt, Fdisp) / res
disp = hs.evaluate_disp(force)
e = hs.evaluate_elastic_energy(force, disp)
kdisp = hs.evaluate_k_disp(force)
ee = hs.evaluate_elastic_energy_k_space(fforce, kdisp)
self.assertTrue(
abs(e - ee) < tol,
"violate Parseval: e = {}, ee = {}, ee/e = {}".format(
e, ee, ee / e))
self.assertTrue(
abs(E - e) < tol, "theoretical E = {}, computed e = {}, "
"diff(tol) = {}({})".format(E, e, E - e, tol))
def real_to_fourier(f, x, y, z):
# Check that real-space grid spacing is all equal
if not (_is_evenly_spaced(x) and _is_evenly_spaced(y) and _is_evenly_spaced(z)):
raise ValueError('Sample points in real space are not evenly spaced.')
dx = x[1]-x[0] # Grid spacing
dy = y[1]-y[0]
dz = z[1]-z[0]
ftrans = ft.rfftn(f)
# Wavenumber arrays
kx = 2*np.pi * ft.fftfreq(x.size, d=dx)
ky = 2*np.pi * ft.fftfreq(y.size, d=dy)
kz = 2*np.pi * ft.rfftfreq(z.size, d=dz) # Only last axis is halved in length when using numpy.fft.rfftn()
# Normalize (convert DFT to continuous FT)
ftrans *= dx*dy*dz
return ftrans, kx, ky, kz
def setup(self, expected_shape):
"""
Setup the convolution with the given shape. The FFT of the PSF will be computed, as well as other small
things that are needed during the convolution step but stay constant if the shape does not change.
:param shape: the shape of the image that will be convoluted
:return: None
"""
self._expected_shape = expected_shape
s1 = array(expected_shape)
s2 = array(self._psf_image.shape)
shape = s1 + s2 - 1
self._fshape = [_next_regular(int(d)) for d in shape]
self._fslice = tuple([slice(0, int(sz)) for sz in shape])
self._psf_fft = rfftn(self._psf_image, self._fshape)
def expand_psf_cube(self):
"""
Interpolate the wavefront pixels to expand the data along lambda
"""
rtf_cube = npf.rfftn(npf.ifftshift(self.psf_cube), axes=(-2, -1))
_, psf_size, rtf_size = rtf_cube.shape
expanded_rtf_cube = np.empty(
(self.expanded_range.size, psf_size, rtf_size), dtype=complex)
for x in range(rtf_size):
for y in range(psf_size):
# Extract the real and imaginary functions at the current point
expanded_pixel = self.expand_carray(rtf_cube[:, y, x],
self.wavelength,
self.expanded_range)
# Store the expanded (x, y) pixel point acorss the full cube.
expanded_rtf_cube[:, y, x] = expanded_pixel
return expanded_rtf_cube
def extended_source_image_(self, ideal_image):
# Convolve
assert np.alltrue(ideal_image.shape == self._expected_shape
), "Shape of image to be convolved is not correct."
ret = irfftn(
rfftn(ideal_image, self._fshape) * self._psf_fft,
self._fshape)[self._fslice].copy()
conv = _centered(ret, self._expected_shape)
#
# fig, sub = plt.subplots(1,1)
#
# #sub[0].imshow(ideal_image, interpolation='none', cmap='gist_heat')
# sub.imshow(conv, interpolation='none', cmap='gist_heat')
#
# fig.savefig("convolution.png")
return conv
def detA_path(A, N=4096):
"""Evaluates det A(lambda) at N equispaced points lambda on interval [0,2pi].
A brief derivation of how this function uses FFT to rapidly evaluate det A(lambda) follows.
We have, letting A_(-j) denote the k*k matrix A[-j,:,:]:
det A(lambda) = det sum_(j=-(T-1))^(T-1) A_(-j)e^(i*j*lambda)
which, flipping the order and realigning j, can be rewritten as
e^(lambda*i*k*(T-1)) det sum_(j=0)^(2T-2) A_(-j+(T-1))e^(-i*j*lambda) (***)
Taking the sum in (***) for the values lambda=0,2*pi/N,...,2*pi*(N-1)/N, assuming N >= (2T-1),
is just taking the discrete Fourier transform of the sequence A_(T-1),...,A_(-(T-1)),0,...,0
right-padded with zeros to length N.
Hence we can rapidly, simultaneously evaluate (***) at all points lambda equispaced from lambda=0
to lambda=2*pi using the FFT. This is implemented below, with additional efficiency from fact that
A(lambda) and A(2*pi-lambda) are conjugate.
"""
# preliminary: assume and verify shape 2*T-1, k, k for A
T = (A.shape[0]+1) // 2
k = A.shape[1]
if not (T == (A.shape[0]+1)/2 and N >= 2*T-1 and k == A.shape[2]):
raise ValueError(f'Asymptotic A matrix has improper shape {A.shape}')
# step 1: use FFT to calculate A(lambda) for each lambda = 2*pi*{0, 1/N, ..., 1/2} (last if N even)
# note that we need to reverse order of A_t to get sequence A_(T-1),...,A_(-(T-1)),0,...,0
Alambda = rfftn(A[::-1,...], axes=(0,), s=(N,))
# step 2: take determinant of each, then multiply by e^(i*k*(T-1)*lambda) to get (***)
det_Alambda = np.empty(N+1, dtype=np.complex128)
det_Alambda[:N//2+1] = np.linalg.det(Alambda)*np.exp(2j*np.pi*k*(T-1)/N*np.arange(N//2+1))
# step 3: use conjugate symmetry to fill in rest
det_Alambda[N//2+1:] = det_Alambda[:(N+1)//2][::-1].conj()
return det_Alambda
def real_to_fourier(func, x, y, z):
# Check that real-space grid spacing is all equal
if not (_is_evenly_spaced(x) and _is_evenly_spaced(y)
and _is_evenly_spaced(z)):
raise ValueError('Sample points in real space are not evenly spaced.')
dx = x[1] - x[0] # Grid spacing
dy = y[1] - y[0]
dz = z[1] - z[0]
ftrans = ft.rfftn(func)
# Wavenumber arrays
kx = 2 * np.pi * ft.fftfreq(x.size, d=dx)
ky = 2 * np.pi * ft.fftfreq(y.size, d=dy)
kz = 2 * np.pi * ft.rfftfreq(
z.size, d=dz
) # Only last axis is halved in length when using numpy.fft.rfftn()
# Normalize (convert DFT to continuous FT)
ftrans *= dx * dy * dz
return ftrans, kx, ky, kz
def fourier_lowpass(x, r):
'''
Fourier lowpass of a 3-D volume
Parameters
==========
x : numpy.ndarray
x.ndim == 3
r : double
Cut-off frequency for Fourier lowpass.
The unit of r is the number of voxels.
The resolution is box_size * voxel_size / r.
'''
fx = rfftn(fftshift(x), norm='ortho')
for (i, j, k) in np.ndindex(fx.shape):
ii = i if i < fx.shape[0] // 2 else i - fx.shape[0]
jj = j if j < fx.shape[1] // 2 else j - fx.shape[1]
kk = k
rho = sqrt(ii**2 + jj**2 + kk**2)
if rho > r:
fx[i, j, k] = 0
return ifftshift(irfftn(fx, norm='ortho'))
def test_realnessEnergy(self):
hs = FreeFFTElasticHalfSpace(self.res, self.young, self.physical_sizes)
force = np.zeros(hs.nb_domain_grid_pts)
force[:self.res[0], :self.res[1]] = np.random.random(self.res)
force[:self.res[0], :self.res[1]] -= \
force[:self.res[0], :self.res[1]].mean()
kdisp = hs.evaluate_k_disp(force)
kforce = rfftn(force.T).T
np_pts = np.prod(hs.nb_domain_grid_pts)
energy = .5 * (np.vdot(-kforce, kdisp) +
np.vdot(-kforce[1:-1, ...], kdisp[1:-1, ...])) / np_pts
error = abs(energy.imag)
tol = 1e-10
self.assertTrue(
error < tol,
"error (imaginary part) = {}, tol = {}".format(error, tol))
error = abs(energy - hs.evaluate_elastic_energy_k_space(kforce, kdisp))
self.assertTrue(error < tol,
("error (comparison) = {}, tol = {}, energy = {}, "
"kenergy = {}").format(
error, tol, energy,
hs.evaluate_elastic_energy_k_space(kforce,
kdisp)))
def _setup_kernel(self):
# voxel indices of array implied by shape
voxels = np.indices(self._bshape, dtype=np.float64)
# coordinates of physical center. XXX - why the 'floor' here?
vox_center = np.floor((np.array(self._bshape) - 1) / 2.0)
phys_center = get_physical_coords(self._affine, vox_center)
# reshape to (N coordinates, -1). We appear to need to assign
# to shape instead of doing a reshape, in order to avoid memory
# copies
voxels.shape = (voxels.shape[0], np.product(voxels.shape[1:]))
# physical coordinates relative to center
X = get_physical_coords(self._affine, voxels) - phys_center
X.shape = (self._ndims[1], ) + tuple(self._bshape)
# compute kernel from these positions
kernel = self(X, axis=0)
kernel = _crop(kernel)
# compute kernel norm
self._norm = _get_kernel_norm(kernel, self._normalization)
self._kernel = kernel
shape_array = (np.ceil(
(np.asarray(self._bshape) + np.asarray(kernel.shape)) / 2) * 2 + 2)
# shape needs to be a list of ints
self._shape = shape_array.astype('uint64').tolist()
self.fkernel = np.zeros(self._shape)
slices = [slice(0, kernel.shape[i]) for i in range(kernel.ndim)]
self.fkernel[slices] = kernel
self.fkernel = npfft.rfftn(self.fkernel)
return kernel
def _cpu_search(self):
d = self.data
c = self.cpu_data
time0 = _time()
for n in xrange(c['rotmat'].shape[0]):
# rotate ligand image
rotate_image3d(c['im_lsurf'], c['vlength'],
np.linalg.inv(c['rotmat'][n]), d['im_center'], c['lsurf'])
c['ft_lsurf'] = rfftn(c['lsurf']).conj()
c['clashvol'] = irfftn(c['ft_lsurf'] * c['ft_rcore'], s=c['shape'])
c['intervol'] = irfftn(c['ft_lsurf'] * c['ft_rsurf'], s=c['shape'])
np.logical_and(c['clashvol'] < c['max_clash'],
c['intervol'] > c['min_interaction'],
c['interspace'])
print('Number of complexes to analyze: ', c['interspace'].sum())
c['chi2'].fill(0)
calc_chi2(c['interspace'], c['q'], c['base_Iq'],
c['rind'], c['rxyz'], c['lind'], (np.mat(c['rotmat'][n])*np.mat(c['lxyz']).T).T,
c['origin'], self.voxelspacing,
c['fifj'], c['targetIq'], c['sq'], c['chi2'])
ind = c['chi2'] > c['best_chi2']
c['best_chi2'][ind] = c['chi2'][ind]
c['rot_ind'][ind] = n
if _stdout.isatty():
self._print_progress(n, c['nrot'], time0)
d['best_chi2'] = c['best_chi2']
d['rot_ind'] = c['rot_ind']
def SmecticInitializer(gridShape, sigma=0.2, seed=None):
if seed is None:
seed = 0
random.seed(seed)
state = SmecticState.SmecticState(gridShape)
field = state.GetOrderParameterField()
Ksq = FourierSpaceTools.FourierSpaceTools(gridShape).kSq.numpy_array()
for component in field.components:
temp = random.normal(scale=gridShape[0],size=gridShape)
ktemp = fft.rfftn(temp)*(sqrt(pi)*sigma)**len(gridShape)*exp(-Ksq*sigma**2/4.)
field[component] = numpy.real(fft.irfftn(ktemp))
## To make seed consistent across grid sizes and convergence comparison
gridShape = copy.copy(oldgrid)
if gridShape[0] != 128:
state = ResizeState(state,gridShape[0],Dim=len(gridShape))
state = ReformatState(state)
state.ktools = FourierSpaceTools.FourierSpaceTools(gridShape)
return state
#Imports
import numpy as np
from numpy.fft import rfftn, irfftn #Importing built in nD Fast Fourier Transforms: rfft2 is nD real fast fourier transform, irfft2 is nD inverse real fast fourier transform
InputArray = np.zeros((100,100,100)) #Array to be fourier transformed
for i in range(100):
for j in range(100):
for k in range(100):
InputArray[i,j,k] = i+j+k
print InputArray[10,10,10]
FT = rfftn(InputArray) #Fourier transform function takes array as input and outputs an array of 1/2 N +1 size. e.g. If InputArray is 100x100x100, FT will be 100x100x51. FT is array of coefficients of sinusoidal functions.
print FT[10,10,10]
#FT array can be manipulated e.g. Removing coefficients for smoothing
IFT = irfftn(FT) #Inverse fourier transform function takes array as input and outputs an array of 2(N-1) size e.g. If FT is 100x100x51, IFT will be 100x100x100
print IFT[10,10,10]
#An explicit Discrete Fourier Transform code is supplied in the slides, but is more complex and less efficient than the in built fast fourier transforms.
def fftconvolve(in1, in2, mode="full", fft_in1=None, fft_in2=None):
"""Convolve two N-dimensional arrays using FFT.
Convolve `in1` and `in2` using the fast Fourier transform method, with
the output size determined by the `mode` argument.
This is generally much faster than `convolve` for large arrays (n > ~500),
but can be slower when only a few output values are needed, and can only
output float arrays (int or object array inputs will be cast to float).
Parameters
----------
in1 : array_like
First input.
in2 : array_like
Second input. Should have the same number of dimensions as `in1`;
if sizes of `in1` and `in2` are not equal then `in1` has to be the
larger array.
mode : str {'full', 'valid', 'same'}, optional
A string indicating the size of the output:
``full``
The output is the full discrete linear convolution
of the inputs. (Default)
``valid``
The output consists only of those elements that do not
rely on the zero-padding.
``same``
The output is the same size as `in1`, centered
with respect to the 'full' output.
Returns
-------
out : array
An N-dimensional array containing a subset of the discrete linear
convolution of `in1` with `in2`.
This code has been modified from scipy.signal.filter.py under the following
license:
Copyright (c) 2001, 2002 Enthought, Inc.
All rights reserved.
Copyright (c) 2003-2012 SciPy Developers.
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
a. Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.
b. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
c. Neither the name of Enthought nor the names of the SciPy Developers
may be used to endorse or promote products derived from this software
without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS
BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
"""
in1 = np.asarray(in1)
in2 = np.asarray(in2)
if np.rank(in1) == np.rank(in2) == 0: # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("in1 and in2 should have the same rank")
elif in1.size == 0 or in2.size == 0: # empty arrays
return np.array([])
s1 = np.array(in1.shape)
s2 = np.array(in2.shape)
complex_result = (np.issubdtype(in1.dtype, np.complex)
or np.issubdtype(in2.dtype, np.complex))
size = s1 + s2 - 1
# Always use 2**n-sized FFT
fsize = 2**np.ceil(np.log2(size)).astype(int)
fslice = tuple([slice(0, int(sz)) for sz in size])
fft1_given = not (fft_in1 is None)
fft2_given = not (fft_in2 is None)
if not complex_result:
if not fft1_given:
fft_in1 = rfftn(in1, fsize)
if not fft2_given:
fft_in2 = rfftn(in2, fsize)
ret = irfftn(fft_in1 * fft_in2, fsize)[fslice].copy()
ret = ret.real
else:
if not fft1_given:
fft_in1 = fftn(in1, fsize)
if not fft2_given:
fft_in2 = fftn(in2, fsize)
ret = ifftn(fft_in1 * fft_in2)[fslice].copy()
if mode == "full":
pass
elif mode == "same":
ret = _centered(ret, s1)
elif mode == "valid":
ret = _centered(ret, s1 - s2 + 1)
else:
raise ValueError("Acceptable mode flags are 'valid',"
" 'same', or 'full'.")
ret = [ret]
if not fft1_given:
ret.append(fft_in1)
if not fft2_given:
ret.append(fft_in2)
ret = tuple(ret)
if len(ret) == 1:
ret = ret[0]
return ret
for i in range(npix+1):
for j in range(npix+1):
for k in range(npix+1):
h = H[i,j,k]
H[ 2*npix-i , j , k ] = h
H[ 2*npix-i , 2*npix-j , k ] = h
H[ 2*npix-i , j , 2*npix-k ] = h
H[ 2*npix-i , 2*npix-j , 2*npix-k ] = h
H[ i , 2*npix-j , k ] = h
H[ i , 2*npix-j , 2*npix-k ] = h
H[ i , j , 2*npix-k ] = h
#print("Progress: %.2f %%\r"%((i/npix)*100), end = '')
#print()
get_H(green_func) # Make the Green's function array
green_ft = rfftn(green_func[:-1,:-1,:-1]) # Take the FT of Green's function array
#plt.imshow(green_func[:,:,0]);plt.colorbar();plt.show()
npix = 2*npix
grid = np.zeros([npix,npix,npix]) # Make the grid array
########################################
### Particle setup ###
########################################
# Positions
x = np.zeros([npt])
y = np.zeros([npt])
z = np.zeros([npt])
def sgimgcoeffs(img, *args, **kwargs):
'''
Given a 3-D image img with shape (nx, ny, nz), use Savitzky-Golay
stencils from savgol(*args, **kwargs) to compute compute the filtered
double-precision image coeffs with shape (nx, ny, nz, ns) such that
coeffs[:,:,:,i] holds the convolution of img with the i-th stencil.
If the image is of single precision, the filter correlation will be done
in single-precision; otherwise, double precision will be used.
The pyfftw module will be used, if available, to accelerate FFT
correlations. Otherwise, the stock Numpy FFT will be used.
'''
# Create the stencils first
stencils = savgol(*args, **kwargs)
if not stencils: raise ValueError('Savitzky-Golay stencil list is empty')
# Make sure the array is in double precision
img = np.asarray(img)
if img.ndim != 3: raise ValueError('Image img must be three-dimensional')
# If possible, find the next-larger efficient size
try: from scipy.fftpack.helper import next_fast_len
except ImportError: next_fast_len = lambda x: x
# Half-sizes of kernels along each axis
hsizes = tuple(bsz // 2 for bsz in stencils[0].shape)
# Padded shape for FFT convolution and the R2C FFT output
pshape = tuple(next_fast_len(isz + 2 * bsz)
for isz, bsz in zip(img.shape, hsizes))
if img.dtype == np.dtype('float32'):
ftype, ctype = np.dtype('float32'), np.dtype('complex64')
else:
ftype, ctype = np.dtype('float64'), np.dtype('complex128')
try:
import pyfftw
except ImportError:
from numpy.fft import rfftn, irfftn
empty = np.empty
use_fftw = False
else:
# Cache PyFFTW planning for 5 seconds
empty = pyfftw.empty_aligned
use_fftw = True
# Build working and output arrays
kernel = empty(pshape, dtype=ftype)
output = empty(img.shape + (len(stencils),), dtype=ftype)
if use_fftw:
# Need to create output arrays and plan both FFTs
krfft = empty(pshape[:-1] + (pshape[-1] // 2 + 1,), dtype=ctype)
rfftn = pyfftw.FFTW(kernel, krfft, axes=(0, 1, 2))
irfftn = pyfftw.FFTW(krfft, kernel,
axes=(0, 1, 2), direction='FFTW_BACKWARD')
m,n,p = img.shape
# Copy the image, leaving space for boundaries
kernel[:,:,:] = 0.
kernel[:m,:n,:p] = img
# For right boundaries, watch for running off left end with small arrays
for ax, (ld, hl) in enumerate(zip(img.shape, hsizes)):
# Build the slice for boundary values
lslices = [slice(None)]*3
rslices = [slice(None)]*3
# Left boundaries are straightforward
lslices[ax] = slice(hl, 0, -1)
rslices[ax] = slice(-hl, None)
kernel[rslices] = kernel[lslices]
# Don't walk off left edge when mirroring right boundary
hi = ld - 1
lo = max(hi - hl, 0)
lslices[ax] = slice(lo, hi)
rslices[ax] = slice(2 * hi - lo, hi, -1)
kernel[rslices] = kernel[lslices]
# Compute the image FFT
if use_fftw:
rfftn.execute()
imfft = krfft.copy()
else: imfft = rfftn(kernel)
i,j,k = hsizes
t,u,v = stencils[0].shape
for l, stencil in enumerate(stencils):
# Clear the kernel storage and copy the stencil
kernel[:,:,:] = 0.
kernel[:t,:u,:v] = stencil[::-1,::-1,::-1]
if use_fftw:
rfftn.execute()
krfft[:,:,:] *= imfft
irfftn(normalise_idft=True)
else: kernel = irfftn(rfftn(kernel) * imfft)
output[:,:,:,l] = kernel[i:i+m,j:j+n,k:k+p]
return output
def _rfftn(a, s=None, axes=None):
return npfft.rfftn(a, s, axes).astype(complex_dtype(a.dtype))
def apply_layer(self, input_image):
# Calculate feed-forward result
assert (input_image.shape[0] == self.ninputs)
if VALID_SIZE_CROP:
# valid size output
output_size = (input_image.shape[1] - self.kernel_size + 1,
input_image.shape[2] - self.kernel_size + 1)
else:
# same size output
output_size = (input_image.shape[1], input_image.shape[2])
output = np.zeros((self.nkernels, output_size[0], output_size[1]),
dtype=np.float32)
self.switches = np.zeros(
(self.nkernels, output_size[0], output_size[1]), dtype=np.uint32)
#options for
#scipy convolution?
#fft convolution?
#cuda convolution?
# Retain precalculated fft / size for efficient repeat calculations
for stridex in range(self.stride_in):
for stridey in range(self.stride_in):
same_fft_size = True
for filteri in range(self.nkernels):
# Apply convolution
if VALID_SIZE_CROP:
stride_shape = (len(
np.arange(
stridex,
input_image.shape[1] - self.kernel_size + 1,
self.stride_in)),
len(
np.arange(
stridey, input_image.shape[2] -
self.kernel_size + 1,
self.stride_in)))
else:
stride_shape = (len(
np.arange(stridex, input_image.shape[1],
self.stride_in)),
len(
np.arange(stridey,
input_image.shape[2],
self.stride_in)))
#conv_result = np.zeros(((output_size[0] + stridex) / self.stride_in, (output_size[1] + stridey) / self.stride_in), dtype=np.float32)
conv_result = np.zeros((stride_shape[0], stride_shape[1]),
dtype=np.float32)
for channeli in range(self.ninputs):
# Space domain convolution
# conv_result = conv_result + convolve2d(
# input_image[channeli, stridex::self.stride_in, stridey::self.stride_in].squeeze(),
# self.W[filteri,channeli,:,:].squeeze(),
# mode='same')
# #mode='valid')
# FFT convolution
#conv_result = conv_result + fftconvolve(
# input_image[channeli, stridex::self.stride_in, stridey::self.stride_in].squeeze(),
# self.W[filteri,channeli,:,:].squeeze(),
# mode='same')
# FFT convolution (cache filter transformations)
convolve_image = input_image[
channeli, stridex::self.stride_in,
stridey::self.stride_in].squeeze()
conv_size = (self.kernel_size +
convolve_image.shape[0] - 1,
self.kernel_size +
convolve_image.shape[1] - 1)
fsize = 2**np.ceil(np.log2(conv_size)).astype(int)
fslice = tuple([slice(0, int(sz)) for sz in conv_size])
if same_fft_size and conv_size == self.prev_conv_size:
fft_result = irfftn(
rfftn(convolve_image, fsize) *
self.Wfft[filteri, channeli, :, :],
fsize)[fslice].copy()
else:
if same_fft_size:
self.Wfft = np.zeros(
(self.nkernels, self.ninputs, fsize[0],
fsize[1] // 2 + 1), np.complex64)
same_fft_size = False
self.prev_conv_size = conv_size
filter_fft = rfftn(
self.W[filteri, channeli, :, :].squeeze(),
fsize)
fft_result = irfftn(
rfftn(convolve_image, fsize) * filter_fft,
fsize)[fslice].copy()
self.Wfft[filteri, channeli, :, :] = filter_fft
conv_result += _centered(fft_result.real,
conv_result.shape)
# if mode == "full":
# return ret
# elif mode == "same":
# return _centered(ret, s1)
# elif mode == "valid":
# return _centered(ret, abs(s1 - s2) + 1)
# Apply maxpool (record switches)
fullx = conv_result.shape[0]
fully = conv_result.shape[1]
splitx = (fullx + 1) / self.maxpool_size
splity = (fully + 1) / self.maxpool_size
striderangex = np.arange(0, fullx - 1, self.maxpool_size)
striderangey = np.arange(0, fully - 1, self.maxpool_size)
for poolx in range(self.maxpool_size):
for pooly in range(self.maxpool_size):
maxpool = np.ones(
(splitx, splity, self.maxpool_size**2),
dtype=np.float32) * -np.inf
offset_i = 0
for offset_x in range(self.maxpool_size):
for offset_y in range(self.maxpool_size):
pool_non_padded = conv_result[
poolx + offset_x::self.maxpool_size,
pooly + offset_y::self.maxpool_size]
maxpool[0:pool_non_padded.shape[0],
0:pool_non_padded.shape[1],
offset_i] = pool_non_padded
offset_i = offset_i + 1
max_indices = np.argmax(maxpool, axis=2)
maxpool = np.amax(maxpool, axis=2)
# Tanh and bias
maxpool = np.tanh(maxpool + self.b[filteri])
# truncate if necessary
if poolx > 0 and fullx % self.maxpool_size >= poolx:
maxpool = maxpool[:-1, :]
max_indices = max_indices[:-1, :]
if pooly > 0 and fully % self.maxpool_size >= pooly:
maxpool = maxpool[:, :-1]
max_indices = max_indices[:, :-1]
output[filteri, stridex +
poolx * self.stride_in::self.stride_out,
stridey + pooly *
self.stride_in::self.stride_out] = maxpool
self.switches[
filteri, stridex +
poolx * self.stride_in::self.stride_out,
stridey + pooly *
self.stride_in::self.stride_out] = max_indices
if filteri == 0:
self.conv_result = conv_result
print "CONV Layer: Done pool {0}, of {1}.".format(
stridex * self.stride_in + stridey + 1, self.stride_in**2)
return output
def _presmooth(self, indata, slices):
_buffer = np.zeros(self._shape)
_buffer[slices] = indata
return npfft.rfftn(_buffer)
def get_ft_image(self):
if self.ft is None:
self.ft = rfftn(self.image, self.imageshape)
return self.ft
def _presmooth(self, indata):
slices = [slice(0, self.bshape[i], 1) for i in range(len(self.shape))]
_buffer = np.zeros(self.shape)
_buffer[slices] = indata
return fft.rfftn(_buffer)
def calculate_6d_integral(self, n_g, q0_g, a2_g=None, e_LDAc_g=None, v_LDAc_g=None, v_g=None, deda2_g=None):
self.timer.start("VdW-DF integral")
self.timer.start("splines")
if self.C_aip is None:
self.construct_cubic_splines()
self.construct_fourier_transformed_kernels()
self.timer.stop("splines")
gd = self.gd
N = self.Nalpha
world = self.world
vdwcomm = self.vdwcomm
if self.alphas:
self.timer.start("hmm1")
i_g = (np.log(q0_g / self.q_a[1] * (self.lambd - 1) + 1) / log(self.lambd)).astype(int)
dq0_g = q0_g - self.q_a[i_g]
self.timer.stop("hmm1")
else:
i_g = None
dq0_g = None
if self.verbose:
print "VDW: fft:",
theta_ak = {}
p_ag = {}
for a in self.alphas:
self.timer.start("hmm2")
C_pg = self.C_aip[a, i_g].transpose((3, 0, 1, 2))
pa_g = C_pg[0] + dq0_g * (C_pg[1] + dq0_g * (C_pg[2] + dq0_g * C_pg[3]))
self.timer.stop("hmm2")
del C_pg
self.timer.start("FFT")
theta_ak[a] = rfftn(n_g * pa_g, self.shape).copy()
if extra_parameters.get("vdw0"):
theta_ak[a][0, 0, 0] = 0.0
self.timer.stop()
if not self.energy_only:
p_ag[a] = pa_g
del pa_g
if self.verbose:
print a,
sys.stdout.flush()
if self.energy_only:
del i_g
del dq0_g
if self.verbose:
print
print "VDW: convolution:",
F_ak = {}
dj_k = self.dj_k
energy = 0.0
for a in range(N):
if vdwcomm is not None:
vdw_ranka = a * vdwcomm.size // N
F_k = np.zeros((self.shape[0], self.shape[1], self.shape[2] // 2 + 1), complex)
self.timer.start("Convolution")
for b in self.alphas:
_gpaw.vdw2(self.phi_aajp[a, b], self.j_k, dj_k, theta_ak[b], F_k)
self.timer.stop()
if vdwcomm is not None:
self.timer.start("gather")
for F in F_k:
vdwcomm.sum(F, vdw_ranka)
# vdwcomm.sum(F_k, vdw_ranka)
self.timer.stop("gather")
if vdwcomm is not None and vdwcomm.rank == vdw_ranka:
if not self.energy_only:
F_ak[a] = F_k
energy += np.vdot(theta_ak[a][:, :, 0], F_k[:, :, 0]).real
energy += np.vdot(theta_ak[a][:, :, -1], F_k[:, :, -1]).real
energy += 2 * np.vdot(theta_ak[a][:, :, 1:-1], F_k[:, :, 1:-1]).real
if self.verbose:
print a,
sys.stdout.flush()
del theta_ak
if self.verbose:
print
if not self.energy_only:
F_ag = {}
for a in self.alphas:
n1, n2, n3 = gd.get_size_of_global_array()
self.timer.start("iFFT")
F_ag[a] = irfftn(F_ak[a]).real[:n1, :n2, :n3].copy()
self.timer.stop()
del F_ak
self.timer.start("potential")
self.calculate_potential(n_g, a2_g, i_g, dq0_g, p_ag, F_ag, e_LDAc_g, v_LDAc_g, v_g, deda2_g)
self.timer.stop()
self.timer.stop()
return 0.5 * world.sum(energy) * gd.dv / self.shape.prod()
def _cpu_search(self):
"""Method which performs the exhaustive search using CPU resources"""
d = self.data
c = self.cpu_data
# initialize the number of total sampled complexes and the number of
# complexes consistent with exactly N restraints
tot_complex = 0
list_total_allowed = np.zeros(max(2, d['nrestraints'] + 1),
dtype=np.float64)
# initalize the time
time0 = _time()
for n in xrange(c['rotmat'].shape[0]):
# rotate the scanning chain object. The rotation needs to be
# inverted, as we are rotating the array, instead of the object.
rotate_image3d(c['im_lsurf'], c['vlength'],
np.linalg.inv(c['rotmat'][n]), d['im_center'],
c['lsurf'])
# calculate the clashing and interaction volume at every position
# in space using FFTs.
np.conj(rfftn(c['lsurf']), c['ft_lsurf'])
c['clashvol'] = irfftn(c['ft_lsurf'] * c['ft_rcore'], s=c['shape'])
c['intervol'] = irfftn(c['ft_lsurf'] * c['ft_rsurf'], s=c['shape'])
# Calculate the accessible interaction space for the current
# rotation. The clashing volume should not be too high, and the
# interaction volume of a reasonable size
np.logical_and(c['clashvol'] < c['max_clash'],
c['intervol'] > c['min_interaction'],
c['interspace'])
# Calculate the number of complexes and multiply with the weight
# for the orientation to correct for rotational/orientational bias
tot_complex += c['weights'][n] * c['interspace'].sum()
# if distance-restraints are available
if self.distance_restraints:
c['restspace'].fill(0)
# determine the center of the distance-restraint consistent
# spheres
rest_center = d['restraints'][:, :3] - \
(np.mat(c['rotmat'][n]) * \
np.mat(d['restraints'][:,3:6]).T).T
mindis = d['restraints'][:, 6]
maxdis = d['restraints'][:, 7]
# Markate the space that is consistent with the distance restraints
distance_restraint(rest_center, mindis, maxdis, c['restspace'])
# Multiply the interaction space with the distance-restraint
# consistent space
c['interspace'] *= c['restspace']
# Now count which violation has been violated
count_violations(rest_center, mindis, maxdis, c['interspace'],
c['weights'][n], c['violations'])
# To visualize the accessible interaction space, keep the maximum
# number of consistent restraints found at every position in space
np.maximum(c['interspace'], c['access_interspace'],
c['access_interspace'])
# Keep track of the number of accessible complexes consistent with
# EXACTLY N restraints. Again, correct for the
# rotational/orientation bias
list_total_allowed += c['weights'][n] *\
np.bincount(c['interspace'].ravel(),
minlength=(max(2, d['nrestraints']+1)))
# Give the user information on progress if it is used interactively
if _stdout.isatty():
self._print_progress(n, c['nrot'], time0)
# attach the output on the self.data dictionary
# the accessible interaction space which will be visualized
d['accessible_interaction_space'] = c['access_interspace']
# the number of accessible complexes consistent with EXACTLY a certain number of restraints
# the number of accessible complexes consistent with EXACTLY a certain
# number of restraints. To account for this, the number of total
# sampled complexes needs to be reduced by the number of complexes
# consistent with 1 or more restraints
d['accessible_complexes'] = [
tot_complex - sum(list_total_allowed[1:])
] + list(list_total_allowed[1:])
# the violation matrix
d['violations'] = c['violations']
def rfft(a, nthreads=ncpu):
return fftw.rfftn(a)
def correlate_windows(window_a,
window_b,
corr_method='fft',
nfftx=None,
nffty=None,
nfftz=None):
"""Compute correlation function between two interrogation windows.
The correlation function can be computed by using the correlation
theorem to speed up the computation.
Parameters
----------
window_a : 2d np.ndarray
a two dimensions array for the first interrogation window,
window_b : 2d np.ndarray
a two dimensions array for the second interrogation window.
corr_method : string
one method is currently implemented: 'fft'.
nfftx : int
the size of the 2D FFT in x-direction,
[default: 2 x windows_a.shape[0] is recommended].
nffty : int
the size of the 2D FFT in y-direction,
[default: 2 x windows_a.shape[1] is recommended].
nfftz : int
the size of the 2D FFT in z-direction,
[default: 2 x windows_a.shape[2] is recommended].
Returns
-------
corr : 3d np.ndarray
a three dimensional array of the correlation function.
Note that due to the wish to use 2^N windows for faster FFT
we use a slightly different convention for the size of the
correlation map. The theory says it is M+N-1, and the
'direct' method gets this size out
the FFT-based method returns M+N size out, where M is the window_size
and N is the search_area_size
It leads to inconsistency of the output
"""
if corr_method == 'fft':
window_b = np.conj(window_b[::-1, ::-1, ::-1])
if nfftx is None:
nfftx = nextpower2(window_b.shape[0] + window_a.shape[0])
if nffty is None:
nffty = nextpower2(window_b.shape[1] + window_a.shape[1])
if nfftz is None:
nfftz = nextpower2(window_b.shape[2] + window_a.shape[2])
f2a = rfftn(normalize_intensity(window_a), s=(nfftx, nffty, nfftz))
f2b = rfftn(normalize_intensity(window_b), s=(nfftx, nffty, nfftz))
corr = irfftn(f2a * f2b).real
corr = corr[:window_a.shape[0] +
window_b.shape[0], :window_b.shape[1] +
window_a.shape[1], :window_b.shape[2] + window_a.shape[2]]
return corr
# elif corr_method == 'direct':
# return convolve2d(normalize_intensity(window_a),
# normalize_intensity(window_b[::-1, ::-1, ::-1]), 'full')
else:
raise ValueError('method is not implemented')
for j in range(dim):
startPoints[i, j] = random.random() * boxSize
(thetas, phis) = hp.pix2ang(healPixResolution,
np.arange(hp.nside2npix(healPixResolution)))
averageDensity = pointMassCount / (linearGridSize)**3 # In count/cell volume
delta_r = gridDensities / averageDensity - 1
k_vect = 2 * math.pi * fft.fftfreq(
linearGridSize, boxSize / linearGridSize) #This is in units of 1/Mpc
k_vectLast = 2 * math.pi * fft.rfftfreq(linearGridSize,
boxSize / linearGridSize)
delta_k = fft.rfftn(delta_r)
k2_inverse = np.empty(delta_k.shape, dtype=np.complex_)
if (dim == 3):
for i in range(delta_k.shape[0]):
for j in range(delta_k.shape[1]):
for k in range(delta_k.shape[2]):
k2_inverse[i, j, k] = 1.0 / (k_vect[i] * k_vect[i] +
k_vect[j] * k_vect[j] +
k_vectLast[k] * k_vectLast[k])
k2_inverse[0, 0, 0] = 0.0
def weightedfftconvolve(in1,
in2,
mode="full",
weighting='none',
displayplots=False):
"""Convolve two N-dimensional arrays using FFT.
Convolve `in1` and `in2` using the fast Fourier transform method, with
the output size determined by the `mode` argument.
This is generally much faster than `convolve` for large arrays (n > ~500),
but can be slower when only a few output values are needed, and can only
output float arrays (int or object array inputs will be cast to float).
Parameters
----------
in1 : array_like
First input.
in2 : array_like
Second input. Should have the same number of dimensions as `in1`;
if sizes of `in1` and `in2` are not equal then `in1` has to be the
larger array.
mode : str {'full', 'valid', 'same'}, optional
A string indicating the size of the output:
``full``
The output is the full discrete linear convolution
of the inputs. (Default)
``valid``
The output consists only of those elements that do not
rely on the zero-padding.
``same``
The output is the same size as `in1`, centered
with respect to the 'full' output.
Returns
-------
out : array
An N-dimensional array containing a subset of the discrete linear
convolution of `in1` with `in2`.
"""
in1 = np.asarray(in1)
in2 = np.asarray(in2)
if np.isscalar(in1) and np.isscalar(in2): # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("in1 and in2 should have the same rank")
elif in1.size == 0 or in2.size == 0: # empty arrays
return np.array([])
s1 = np.array(in1.shape)
s2 = np.array(in2.shape)
complex_result = (np.issubdtype(in1.dtype, np.complex)
or np.issubdtype(in2.dtype, np.complex))
size = s1 + s2 - 1
if mode == "valid":
_check_valid_mode_shapes(s1, s2)
# Always use 2**n-sized FFT
fsize = 2**np.ceil(np.log2(size)).astype(int)
fslice = tuple([slice(0, int(sz)) for sz in size])
if not complex_result:
fft1 = rfftn(in1, fsize)
fft2 = rfftn(in2, fsize)
theorigmax = np.max(
np.absolute(irfftn(gccproduct(fft1, fft2, 'none'), fsize)[fslice]))
ret = irfftn(
gccproduct(fft1, fft2, weighting, displayplots=displayplots),
fsize)[fslice].copy()
ret = irfftn(
gccproduct(fft1, fft2, weighting, displayplots=displayplots),
fsize)[fslice].copy()
ret = ret.real
ret *= theorigmax / np.max(np.absolute(ret))
else:
fft1 = fftpack.fftn(in1, fsize)
fft2 = fftpack.fftn(in2, fsize)
theorigmax = np.max(
np.absolute(fftpack.ifftn(gccproduct(fft1, fft2, 'none'))[fslice]))
ret = fftpack.ifftn(
gccproduct(fft1, fft2, weighting,
displayplots=displayplots))[fslice].copy()
ret *= theorigmax / np.max(np.absolute(ret))
# scale to preserve the maximum
if mode == "full":
return ret
elif mode == "same":
return _centered(ret, s1)
elif mode == "valid":
return _centered(ret, s1 - s2 + 1)
if __name__ == "__main__":
if torch.cuda.is_available():
nfft3 = lambda x: nfft.fftn(x, axes=(1, 2, 3))
nifft3 = lambda x: nfft.ifftn(x, axes=(1, 2, 3))
cfs = [cfft.fft, cfft.fft2, cfft.fft3]
nfs = [nfft.fft, nfft.fft2, nfft3]
cifs = [cfft.ifft, cfft.ifft2, cfft.ifft3]
nifs = [nfft.ifft, nfft.ifft2, nifft3]
for args in zip(cfs, nfs, cifs, nifs):
test_c2c(*args)
nrfft3 = lambda x: nfft.rfftn(x, axes=(1, 2, 3))
nirfft3 = lambda x: nfft.irfftn(x, axes=(1, 2, 3))
cfs = [cfft.rfft, cfft.rfft2, cfft.rfft3]
nfs = [nfft.rfft, nfft.rfft2, nrfft3]
cifs = [cfft.irfft, cfft.irfft2, cfft.irfft3]
nifs = [nfft.irfft, nfft.irfft2, nirfft3]
for args in zip(cfs, nfs, cifs, nifs):
test_r2c(*args)
test_expand()
test_fft_gradcheck()
test_ifft_gradcheck()
test_fft2d_gradcheck()
test_ifft2d_gradcheck()
def fftconvolve_new(in1, in2, mode="full"):
"""Convolve two N-dimensional arrays using FFT.
Convolve `in1` and `in2` using the fast Fourier transform method, with
the output size determined by the `mode` argument.
This is generally much faster than `convolve` for large arrays (n > ~500),
but can be slower when only a few output values are needed, and can only
output float arrays (int or object array inputs will be cast to float).
Parameters
----------
in1 : array_like
First input.
in2 : array_like
Second input. Should have the same number of dimensions as `in1`;from scipy.signal import fftconvolve
if sizes of `in1` and `in2` are not equal then `in1` has to be the
larger array.get_window
mode : str {'full', 'valid', 'same'}, optional
A string indicating the size of the output:
``full``
The output is the full discrete linear convolution
of the inputs. (Default)
``valid``
The output consists only of those elements that do not
rely on the zero-padding.
``same``
The output is the same size as `in1`, centered
with respect to the 'full' output.
Returns
-------
out : array
An N-dimensional array containing a subset of the discrete linear
convolution of `in1` with `in2`.
Examples
--------
Autocorrelation of white noise is an impulse. (This is at least 100 times
as fast as `convolve`.)
>>> from scipy import signal
>>> sig = np.random.randn(1000)
>>> autocorr = signal.fftconvolve(sig, sig[::-1], mode='full')
>>> import matplotlib.pyplot as plt
>>> fig, (ax_orig, ax_mag) = plt.subplots(2, 1)
>>> ax_orig.plot(sig)
>>> ax_orig.set_title('White noise')
>>> ax_mag.plot(np.arange(-len(sig)+1,len(sig)), autocorr)
>>> ax_mag.set_title('Autocorrelation')
>>> fig.tight_layout()
>>> fig.show()
Gaussian blur implemented using FFT convolution. Notice the dark borders
around the image, due to the zero-padding beyond its boundaries.
The `convolve2d` function allows for other types of image boundaries,
but is far slower.
>>> from scipy import misc
>>> lena = misc.lena()
>>> kernel = np.outer(signal.gaussian(70, 8), signal.gaussian(70, 8))
>>> blurred = signal.fftconvolve(lena, kernel, mode='same')
>>> fig, (ax_orig, ax_kernel, ax_blurred) = plt.subplots(1, 3)
>>> ax_orig.imshow(lena, cmap='gray')
>>> ax_orig.set_title('Original')
>>> ax_orig.set_axis_off()
>>> ax_kernel.imshow(kernel, cmap='gray')
>>> ax_kernel.set_title('Gaussian kernel')
>>> ax_kernel.set_axis_off()
>>> ax_blurred.imshow(blurred, cmap='gray')
>>> ax_blurred.set_title('Blurred')
>>> ax_blurred.set_axis_off()
>>> fig.show()
"""
in1 = asarray(in1)
in2 = asarray(in2)
if in1.ndim == in2.ndim == 0: # scalar inputs
return in1 * in2
elif not in1.ndim == in2.ndim:
raise ValueError("in1 and in2 should have the same dimensionality")
elif in1.size == 0 or in2.size == 0: # empty arrays
return array([])
s1 = array(in1.shape)
s2 = array(in2.shape)
complex_result = (np.issubdtype(in1.dtype, np.complex)
or np.issubdtype(in2.dtype, np.complex))
shape = s1 + s2 - 1
if mode == "valid":
_check_valid_mode_shapes(s1, s2)
# Speed up FFT by padding to optimal size for FFTPACK
# expand by at least twice+1
fshape = [_next_regular(int(d)) for d in shape]
fslice = tuple([slice(0, int(sz)) for sz in shape])
# Pre-1.9 NumPy FFT routines are not threadsafe. For older NumPys, make
# sure we only call rfftn/irfftn from one thread at a time.
if not complex_result and (_rfft_mt_safe or _rfft_lock.acquire(False)):
try:
ret = irfftn(rfftn(in1, fshape) * rfftn(in2, fshape),
fshape)[fslice].copy()
finally:
if not _rfft_mt_safe:
_rfft_lock.release()
else:
# If we're here, it's either because we need a complex result, or we
# failed to acquire _rfft_lock (meaning rfftn isn't threadsafe and
# is already in use by another thread). In either case, use the
# (threadsafe but slower) SciPy complex-FFT routines instead.
ret = ifftn(fftn(in1, fshape) * fftn(in2, fshape))[fslice].copy()
if not complex_result:
ret = ret.real
if mode == "full":
return ret
elif mode == "same":
return _centered(ret, s1)
elif mode == "valid":
return _centered(ret, s1 - s2 + 1)
else:
raise ValueError("Acceptable mode flags are 'valid',"
" 'same', or 'full'.")
meanx, meany = np.random.uniform(padding + 2.,
shape[0] - padding - 2.,
size=2)
foo = -0.5 * ((xx - meanx)**2 + (yy - meany)**2) / sigma2
trueimage += np.exp(foo)
for i in range(10):
x1, y1 = np.random.uniform(padding + 1.,
shape[0] - padding - 7.,
size=2)
dx1, dy1 = np.random.uniform(1., 6., size=2)
trueimage[y1:y1 + dy1, x1:x1 + dx1] += 0.5
trueimage[:padding, :] = 0.
trueimage[:, :padding] = 0.
trueimage[-padding:, :] = 0.
trueimage[:, -padding:] = 0.
trueft = rfftn(trueimage, shape)
data = (trueft * trueft.conj()).real
# construct an inverse variance "noise level"
sigma = np.zeros_like(data) + 0.05 * np.median(data)
sigma2 += sigma**2 + (0.05 * data)**2
ivar = 1. / sigma2
# construct and test class
model = pharetModel(data, shape, padding, ivar=ivar)
# initialize emcee
ndim = 32 * 32
nwalkers = 2 * ndim + 2
pos = np.random.normal(size=(nwalkers, ndim))
sampler = emcee.EnsembleSampler(nwalkers, ndim, model, args=[
def invert_jacdict(jacdict, unknowns, targets, tau, test_invertible=False):
"""Given a nested dict of ATI Jacobians that maps unknowns -> targets, e.g. an asymptotic
H_U matrix, get the inverse H_U^(-1) as a nested dict.
This is implemented by inverting the FFT-based multiplication that was implemented above
for ATI, making use of the linearity of the FFT:
- We take the FFT of each ATI Jacobian, padded out to 4*tau-3 as above
(This is done by first packing all Jacobians into a single array A)
- Then, we take the FFT of the identity, centered aroun d2*tau-1 since
we intend it to be the result of a product
- We solve frequency-by-frequency, i.e. for each of 4*tau-3 omegas we solve a k*k
linear system to get A_rfft[omega,...]^(-1)*id_rfft[omega,...]
- We take the inverse FFT of the results, then take only the first 2*tau-1 elements
to get (approximate) inverse Jacobians with times -(tau-1),...,(tau-1), same as
original Jacobians
- We unpack these to get a nested dict of ATI Jacobians that inverts original 'jacdict'
Parameters
----------
jacdict : dict of dict, ATI (or convertible to ATI) Jacobians where jacdict[t][u] gives
asymptotic mapping from unknowns u to targets t in H_U
unknowns : list, names of unknowns in H_U
targets : list, names of targets in H_U
tau : int, convert all ATI Jacobians to size tau and provide inverse in size tau
test_invertible : [optional] bool, use winding number criterion to test whether we should
really be inverting this system (i.e. whether determinate solution)
Returns
-------
inv_jacdict : dict of dict, ATI Jacobians where inv_jacdict[u][t] gives asymptotic mapping
from targets t to unknowns u in H_U^(-1)
"""
k = len(unknowns)
assert k == len(targets)
# stack the k^2 Jacobians relating unknowns to targets into an A matrix
A = jac.pack_asymptotic_jacobians(jacdict, unknowns, targets, tau)
if test_invertible:
# use winding number criterion to test invertibility
if determinacy.winding_criterion(A, N=4096) != 0:
raise ValueError('Trying to invert asymptotic time invariant system of Jacobians' +
' but winding number test says that it is not uniquely invertible!')
# take FFT of first dimension (time) of A (i.e. take FFT separtely of all k^2 Jacobians)
A_rfft = rfftn(A, s=(4*tau-3,), axes=(0,))
# take FFT of identity operator (for efficiency, reuse smaller calc)
id_vec_rfft = rfft(np.arange(4*tau-3)==(2*tau-2))
id_rfft = np.zeros((2*tau-1, k, k), dtype=np.complex128)
for i in range(k):
id_rfft[:, i, i] = id_vec_rfft
# now solve the linear system to invert A frequency-by-frequency
# (since frequency is leading dimension, np.linalg.solve automatically does this)
A_rfft_inv = np.linalg.solve(A_rfft, id_rfft)
# take inverse FFT of this to get full A
# then take first 2*tau-1 entries to get approximate A from -(tau-1),...,0,...,(tau-1)
A_inv = irfftn(A_rfft_inv, s=(4*tau-3,), axes=(0,))[:2*tau-1, :, :]
# unstack this
return jac.unpack_asymptotic_jacobians(A_inv, targets, unknowns, tau)
def test_r2c_outofplace(self):
"""
Test out-of-place R2C transforms
"""
n = 32
for dims in range(1, 5):
if dims >= 3:
ndim_max = min(dims + 1, 2)
else:
ndim_max = min(dims + 1, 3)
for ndim in range(1, ndim_max):
for dtype in [np.float32, np.float64]:
for norm in [0, 1, "ortho"]:
with self.subTest(dims=dims,
ndim=ndim,
dtype=dtype,
norm=norm):
if dtype == np.float32:
rtol = 1e-6
else:
rtol = 1e-12
if dtype == np.float32:
dtype_c = np.complex64
elif dtype == np.float64:
dtype_c = np.complex128
sh = [n] * dims
sh = tuple(sh)
shc = [n] * dims
shc[-1] = n // 2 + 1
shc = tuple(shc)
d = np.random.uniform(0, 1, sh).astype(dtype)
# A pure random array may not be a very good test (too random),
# so add a Gaussian
xx = [
np.fft.fftshift(np.fft.fftfreq(nn))
for nn in sh
]
v = np.zeros_like(d)
for x in np.meshgrid(*xx, indexing='ij'):
v += x**2
d += 10 * np.exp(-v * 2)
n0 = (abs(d)**2).sum()
d_gpu = cua.to_gpu(d)
d1_gpu = cua.empty(shc, dtype=dtype_c)
app = VkFFTApp(d.shape,
d.dtype,
ndim=ndim,
norm=norm,
r2c=True,
inplace=False)
# base FFT scale
s = np.sqrt(np.prod(d.shape[-ndim:]))
d = rfftn(d, axes=list(range(dims))[-ndim:]) / s
d1_gpu = app.fft(d_gpu, d1_gpu)
d1_gpu *= dtype_c(app.get_fft_scale())
self.assertTrue(d1_gpu.shape == tuple(shc))
self.assertTrue(d1_gpu.dtype == dtype_c)
self.assertTrue(
np.allclose(d,
d1_gpu.get(),
rtol=rtol,
atol=abs(d).max() * rtol))
d = irfftn(d, axes=list(range(dims))[-ndim:]) * s
d_gpu = app.ifft(d1_gpu, d_gpu)
d_gpu *= dtype(app.get_ifft_scale())
self.assertTrue(d_gpu.shape == tuple(sh))
self.assertTrue(
np.allclose(d,
d_gpu.get(),
rtol=rtol,
atol=abs(d).max() * rtol))
n1 = (abs(d_gpu.get())**2).sum()
self.assertTrue(np.isclose(n0, n1, rtol=rtol))
def fft(x, ax, ncpu):
return rfftn(x, axes=ax)
