def GetIFFT3Dfield(self, cv):
nz_half = self.nz // 2
temp_cv = ft.zeros_aligned((self.nx, self.ny, nz_half + 1),
dtype='complex64')
v = ft.zeros_aligned((self.nx, self.ny, self.nz), dtype='float32')
# The FFTW instance kills the input array, so don't pass the actual array here.
fftyb = ft.FFTW(temp_cv,
temp_cv,
direction='FFTW_BACKWARD',
axes=(1, ),
flags=('FFTW_MEASURE', ))
rfftzb = ft.FFTW(temp_cv,
v,
direction='FFTW_BACKWARD',
axes=(2, ),
flags=('FFTW_MEASURE', ))
cvt = ft.zeros_aligned((self.nx, self.ny, nz_half + 1),
dtype='complex64')
np.copyto(cvt, cv)
fftyb.update_arrays(cvt, cvt)
fftyb.execute()
self.DoMPITranspose(cvt, self.nx, nz_half, self.nproc, self.my_id)
fftyb.execute()
rfftzb.update_arrays(cvt, v)
rfftzb.execute()
return v
def GetFFT3Dfield(self, v):
nz_half = self.nz // 2
cv = ft.zeros_aligned((self.nx, self.ny, nz_half + 1),
dtype='complex64')
temp_v = ft.zeros_aligned((self.nx, self.ny, self.nz), dtype='float32')
# The FFTW instance kills the input array, so don't pass the actual array here.
rfftzf = ft.FFTW(temp_v,
cv,
direction='FFTW_FORWARD',
axes=(2, ),
flags=('FFTW_MEASURE', ))
fftyf = ft.FFTW(cv,
cv,
direction='FFTW_FORWARD',
axes=(1, ),
flags=('FFTW_MEASURE', ))
rfftzf.update_arrays(v, cv)
rfftzf.execute()
fftyf.update_arrays(cv, cv)
fftyf.execute()
self.DoMPITranspose(cv, self.nx, nz_half, self.nproc, self.my_id)
fftyf.execute()
alpha2 = 1.0 / (float(self.ny) * float(self.ny) * float(self.ny))
cv[:, :, :] = cv[:, :, :] * alpha2
return cv
def __init__(self,lx,ly,lz,nproc,my_id):
self.lx = lx
self.ly = ly
self.lz = lz
self.nproc = nproc
self.my_id = my_id
#self.transpose=MPITranspose()
self.transpose=MPITranspose(lx,ly,lz)
self.temp_v = ft.zeros_aligned((lx,ly,lz), dtype='float32')
self.temp_cv = ft.zeros_aligned((lx,ly,1+(lz//2)), dtype='complex64')
self.cv = ft.zeros_aligned((lx,ly,1+(lz//2)), dtype='complex64')
self.v = ft.zeros_aligned((lx,ly,lz), dtype='float32')
self.cvt = ft.zeros_aligned((lx,ly,1+(lz//2)), dtype='complex64')
self.rfftzf = ft.FFTW(self.temp_v,self.cv,direction='FFTW_FORWARD',axes=(2,),flags=('FFTW_MEASURE',))
self.fftyf = ft.FFTW(self.cv,self.cv,direction='FFTW_FORWARD',axes=(1,),flags=('FFTW_MEASURE',))
self.ifftyb = ft.FFTW(self.temp_cv,self.temp_cv,direction='FFTW_BACKWARD',axes=(1,),flags=('FFTW_MEASURE',))
self.irfftzb = ft.FFTW(self.temp_cv,self.v,direction='FFTW_BACKWARD',axes=(2,),flags=('FFTW_MEASURE',))
self.alpha2 = 1.0/(float(ly)*float(ly)*float(lz))
return
def DownldData_cSOAP(self, dataset_name, time, lx, ly, lz, nproc, my_id,
auth_token):
vx = ft.zeros_aligned((lx, ly, lz), dtype='float32')
vy = ft.zeros_aligned((lx, ly, lz), dtype='float32')
vz = ft.zeros_aligned((lx, ly, lz), dtype='float32')
SOAPtdb.loadvel(vx, vy, vz, lx, ly, lz, my_id)
return vx, vy, vz
def DownldData_pyJHTDB(self,dataset_name,time,lx,ly,lz,nproc,my_id,auth_token,getFunction='Velocity'):
chkSz=32 #This is the maximum possible. May be increased in future depending on network bandwidth
slabs=lx//chkSz
lJHTDB=libJHTDB(auth_token)
lJHTDB.initialize()#NOTE: datbase returns Velcocity as [lz,ly,lx,3]
for k in range(slabs):
start=np.array([my_id*lx+k*chkSz,0,0],dtype=np.int)
width=np.array([chkSz,ly,lz],dtype=np.int)
uAll=lJHTDB.getRawData(time,start,width,data_set=dataset_name,getFunction=getFunction)
if(k==0):
vx=uAll[:,:,:,0]
vy=uAll[:,:,:,1]
vz=uAll[:,:,:,2]
else:
vx=np.concatenate((vx,uAll[:,:,:,0]),axis=2) #axis=2=> the index of lx
vy=np.concatenate((vy,uAll[:,:,:,1]),axis=2)
vz=np.concatenate((vz,uAll[:,:,:,2]),axis=2)
lJHTDB.finalize()
u=ft.zeros_aligned((lx,ly,lz),dtype='float32')
v=ft.zeros_aligned((lx,ly,lz),dtype='float32')
w=ft.zeros_aligned((lx,ly,lz),dtype='float32')
u[:,:,:]=np.transpose(vx)
v[:,:,:]=np.transpose(vy)
w[:,:,:]=np.transpose(vz)
return u,v,w
def get_matrices(M, n):
""" Make sure M is aligned and generate n other matrices. """
print "Aligning to ", simd_alignment, " bytes"
m = zeros_aligned(M.shape)
np.copyto(m, M)
M = m
lst = []
for _ in xrange(n):
lst.append(zeros_aligned(M.shape))
lst.append(M)
return lst
def _cwt_full_fftw(data_fft, wavelet, scale, derivative, cwt, ii, extend_len,
threads, omega):
out = pyfftw.zeros_aligned((1, len(data_fft)), dtype='complex128')
fft_conv_inv = pyfftw.FFTW(out,
out,
axes=(-1, ),
direction='FFTW_BACKWARD',
flags=['FFTW_ESTIMATE'],
threads=threads)
wavelet.freq_domain_numba(omega,
np.atleast_1d(scale),
out,
derivative=derivative)
if not wavelet.is_analytic:
out[:] = out.conj()
out *= data_fft
fft_conv_inv(normalise_idft=True)
cwt[ii:ii + 1] = out[:, extend_len:len(data_fft) - extend_len]
return
def check_array(self, array, shape, dtype, copy=True):
"""
Check that a given array is compatible with the FFTW plans,
in terms of alignment and data type.
If the provided array does not meet any of the checks, a new array
is returned.
"""
if array.shape != shape:
raise ValueError("Invalid data shape: expected %s, got %s" %
(shape, array.shape)
)
if array.dtype != dtype:
raise ValueError("Invalid data type: expected %s, got %s" %
(dtype, array.dtype)
)
if self.check_alignment and not(pyfftw.is_byte_aligned(array)):
array2 = pyfftw.zeros_aligned(self.shape, dtype=self.dtype_in)
np.copyto(array2, array)
else:
if copy:
array2 = np.copy(array)
else:
array2 = array
return array2
def coarseRicci(L, sqdist, R, temp1=None, temp2=None):
"""
Fully optimized Ricci matrix computation.
Requires 7 matrix multiplications and many entrywise operations.
Only 2 temporary matrices are needed, and can be provided as arguments.
Uses full gemm functionality to avoid creating intermediate matrices.
R is the output array, while temp1 and temp2 are temporary matrices.
"""
D = sqdist
if temp1 is None:
temp1 = zeros_aligned(sqdist.shape, n=32)
if temp2 is None:
temp2 = zeros_aligned(sqdist.shape, n=32)
A = temp1
B = temp2
# this C should not exist
B = ne.evaluate("D*D/4.0")
L.dot(B, out=A)
L.dot(D, out=B)
ne.evaluate("A-D*B", out=A)
L.dot(A, out=R)
# the first two terms done
L.dot(B, out=A)
ne.evaluate("R+0.5*(D*A+B*B)", out=R)
# Now R contains everything under overline
ne.evaluate("R+dR-0.5*dA*D-dB*B",
global_dict={
'dA': np.diag(A).copy()[:, None],
'dB': np.diag(B).copy()[:, None],
'dR': np.diag(R).copy()[:, None]
},
out=R)
# Now R contains all but two matrix products from line 2
L.dot(L, out=A)
ne.evaluate("L*BT-0.5*A*D", global_dict={'BT': B.T}, out=A)
add_AB_to_C(A, D, R)
ne.evaluate("L*D", out=A)
add_AB_to_C(A, B, R)
# done!
np.fill_diagonal(R, 0.0)
def calculate_energies(
save_options,
resol,
psi,
cmass,
distarray,
Vcell,
phisp,
karray2,
funct,
fft_psi,
ifft_funct,
egpcmlist,
egpsilist,
ekandqlist,
egylist,
mtotlist,
scalefactor,
scalefactorlist,
):
if (save_options[3]):
egyarr = pyfftw.zeros_aligned((resol, resol, resol), dtype='float64')
# Gravitational potential energy density associated with the central potential
egyarr = ne.evaluate('real((abs(psi))**2)') #compute physical density
egyarr = ne.evaluate('real(-cmass/(distarray*scalefactor)*egyarr)'
) #compute physical energy density
egpcmlist.append(
Vcell *
np.sum(egyarr)) #add total gravitational energy to the list
tot = Vcell * np.sum(
egyarr) #add total gravitational energy to the total energy
# Gravitational potential energy density of self-interaction of the condensate
egyarr = ne.evaluate('real(0.5*(phisp)*real((abs(psi))**2))')
egpsilist.append(Vcell * np.sum(egyarr))
tot = tot + Vcell * np.sum(egyarr)
funct = fft_psi(psi)
funct = ne.evaluate('-karray2*scalefactor**(-2.)*funct')
funct = ifft_funct(funct)
egyarr = ne.evaluate('real((-0.5*conj(psi)*funct))')
ekandqlist.append(Vcell * np.sum(egyarr))
tot = tot + Vcell * np.sum(egyarr)
egylist.append(tot) #Physical total energy
egyarr = ne.evaluate('real(((abs(psi))**2)/scalefactor**3)')
mtotlist.append(Vcell * scalefactor**3 * np.sum(egyarr))
scalefactorlist.append(scalefactor)
def coarseRicci(L, sqdist, R, temp1=None, temp2=None):
"""
Fully optimized Ricci matrix computation.
Requires 7 matrix multiplications and many entrywise operations.
Only 2 temporary matrices are needed, and can be provided as arguments.
Uses full gemm functionality to avoid creating intermediate matrices.
R is the output array, while temp1 and temp2 are temporary matrices.
"""
D = sqdist
if temp1 is None:
temp1 = zeros_aligned(sqdist.shape, n=32)
if temp2 is None:
temp2 = zeros_aligned(sqdist.shape, n=32)
A = temp1
B = temp2
# this C should not exist
B = ne.evaluate("D*D/4.0")
L.dot(B, out=A)
L.dot(D, out=B)
ne.evaluate("A-D*B", out=A)
L.dot(A, out=R)
# the first two terms done
L.dot(B, out=A)
ne.evaluate("R+0.5*(D*A+B*B)", out=R)
# Now R contains everything under overline
ne.evaluate("R+dR-0.5*dA*D-dB*B",
global_dict={'dA': np.diag(A).copy()[:, None],
'dB': np.diag(B).copy()[:, None],
'dR': np.diag(R).copy()[:, None]}, out=R)
# Now R contains all but two matrix products from line 2
L.dot(L, out=A)
ne.evaluate("L*BT-0.5*A*D", global_dict={'BT': B.T}, out=A)
add_AB_to_C(A, D, R)
ne.evaluate("L*D", out=A)
add_AB_to_C(A, B, R)
# done!
np.fill_diagonal(R, 0.0)
def plan_fft(frames_init, dims1, prealloc=True):
'''Plan FFT for pyfftw for a 3D video.
Inputs:
frames_init (int): Number of images.
dims1 (tuplel of int, shape = (2,)): lateral dimension of the image.
prealloc (bool, default to True): True if pre-allocate memory space for large variables.
Achieve faster speed at the cost of higher memory occupation.
Outputs:
bb(3D numpy.ndarray of float32): array of the real video.
bf(3D numpy.ndarray of complex64): array of the complex spectrum.
fft_object_b(pyfftw.FFTW): Object for forward FFT.
fft_object_c(pyfftw.FFTW): Object for inverse FFT.
'''
(rows1, cols1) = dims1
bb = pyfftw.zeros_aligned((frames_init, rows1, cols1),
dtype='float32',
n=8)
if prealloc:
bf = pyfftw.zeros_aligned((frames_init, rows1, cols1 // 2 + 1),
dtype='complex64',
n=8)
else:
bf = pyfftw.empty_aligned((frames_init, rows1, cols1 // 2 + 1),
dtype='complex64',
n=8)
fft_object_b = pyfftw.FFTW(bb,
bf,
axes=(-2, -1),
flags=('FFTW_MEASURE', ),
direction='FFTW_FORWARD',
threads=mp.cpu_count())
fft_object_c = pyfftw.FFTW(bf,
bb,
axes=(-2, -1),
flags=('FFTW_MEASURE', ),
direction='FFTW_BACKWARD',
threads=mp.cpu_count())
return bb, bf, fft_object_b, fft_object_c
def test_zeros_aligned(self):
shape = (10, 10)
# Test a few alignments and dtypes
for each in [(3, 'float64'), (7, 'float64'), (9, 'float32'),
(16, 'int64'), (24, 'bool'), (23, 'complex64'),
(63, 'complex128'), (64, 'int8')]:
n = each[0]
a = numpy.zeros(shape, dtype=each[1])
b = zeros_aligned(shape, dtype=each[1], n=n)
self.assertTrue(b.ctypes.data % n == 0)
self.assertTrue(b.dtype == each[1])
self.assertTrue(numpy.array_equal(a, b))
def updatePSF(self, psf):
"""Set convolution code to use new PSF."""
self.psf = psf
imgshape = self.psf.shape
self.in_realimg = pyfftw.zeros_aligned(imgshape, dtype=N.float32)
self.in_realpsf = pyfftw.zeros_aligned(imgshape, dtype=N.float32)
# output real->complex
cshape = (imgshape[0], imgshape[1] // 2 + 1)
self.temp_cmplx = pyfftw.zeros_aligned(cshape, dtype=N.complex64)
self.temp_cmplxpsf = pyfftw.zeros_aligned(cshape, dtype=N.complex64)
self.out_real = pyfftw.zeros_aligned(imgshape, dtype=N.float32)
self.fwd_fftimg = pyfftw.FFTW(
self.in_realimg,
self.temp_cmplx,
axes=(0, 1),
direction='FFTW_FORWARD',
flags=['FFTW_MEASURE'],
)
self.bkd_fftout = pyfftw.FFTW(
self.temp_cmplx,
self.out_real,
axes=(0, 1),
direction='FFTW_BACKWARD',
flags=['FFTW_MEASURE'],
)
# compute fft of psf (only used once)
self.fwd_fftpsf = pyfftw.FFTW(
self.in_realpsf,
self.temp_cmplxpsf,
axes=(0, 1),
direction='FFTW_FORWARD',
flags=['FFTW_ESTIMATE'],
)
self.in_realpsf[:, :] = psf
self.fwd_fftpsf()
def metricize_pureC(dist, temp=None, limit=0):
"""
Metricize based on BLIS framework for BLAS.
Modified ulmBLAS code for dgemm_nn.
"""
if temp is None:
temp = zeros_aligned(dist.shape, n=32)
# We use subtropical matrix multiplication since it is faster
# Starting with Skylake the tropical one will be as fast
ne.evaluate('exp(sqrt(dist))', out=dist)
np.copyto(temp, dist)
metricize_gemm_pureC(dist, temp, limit)
ne.evaluate('log(dist)**2', out=dist)
def __init__(self, img, width=8, step=4, gamma=3.0, beta_percentile=50, beta_count=0, types='float', mode='numpy', debug=False):
# tried to use pyfftw, but it was slower than numpy fft
import pyfftw
import cv2
# prepare variables
extended_img = cv2.copyMakeBorder(img, 2*width, 2*width, 2*width, 2*width, cv2.BORDER_REPLICATE)
self._img = extended_img
self._xwidth, self._ywidth = width, width
self._xstep, self._ystep = step, step
self._types = types
self._debug = debug
#self._nthread = multiprocessing.cpu_count()
self._nthread = 1
self.NX, self.NY = 2 * self._xwidth + 1, 2 * self._ywidth + 1
self.image_section = pyfftw.zeros_aligned((self.NX, self.NY), dtype='float32')
self.fourier_section = pyfftw.zeros_aligned((self.NX, self.NY), dtype='complex64')
self.method = mode
self.gamma = gamma
self.coords = self._build_coordinates()
self.coord_N = len(self.coords)
self.hanning = _build_hanning_window(self.NX, self.NY)
if beta_count == 0:
beta_count = int(self.coord_N * 1.0)
self.beta_count = beta_count
self.beta_percentile = beta_percentile
self.beta = []
pyfftw.forget_wisdom()
def metricize_mul(dist, temp=None, limit=0):
"""
Metricize based on BLIS framework for BLAS.
Modified ulmBLAS code for dgemm_nn with optimal kernel.
If limit is larger than 0, then only this many rounds will happen.
"""
if temp is None:
temp = zeros_aligned(dist.shape)
# We use subtropical matrix multiplication since it is faster
# Starting with Skylake the tropical one will be as fast
ne.evaluate('exp(sqrt(dist))', out=dist)
np.copyto(temp, dist)
metricize_gemm(dist, temp, limit=4)
ne.evaluate('log(dist)**2', out=dist)
def calculate_energies(
save_options,
resol,
psi,
cmass,
distarray,
Vcell,
phisp,
karray2,
funct,
fft_psi,
ifft_funct,
egpcmlist,
egpsilist,
ekandqlist,
egylist,
mtotlist,
):
if (save_options[3]):
egyarr = pyfftw.zeros_aligned((resol, resol, resol), dtype='float64')
# Gravitational potential energy density associated with the central potential
egyarr = ne.evaluate('real((abs(psi))**2)')
egyarr = ne.evaluate('real((-cmass/distarray)*egyarr)')
egpcmlist.append(Vcell * np.sum(egyarr))
tot = Vcell * np.sum(egyarr)
# Gravitational potential energy density of self-interaction of the condensate
egyarr = ne.evaluate(
'real(0.5*(phisp+(cmass)/distarray)*real((abs(psi))**2))')
egpsilist.append(Vcell * np.sum(egyarr))
tot = tot + Vcell * np.sum(egyarr)
# TODO: Does this reuse the memory of funct? That is the
# intention, but likely isn't what is happening
funct = fft_psi(psi)
funct = ne.evaluate('-karray2*funct')
funct = ifft_funct(funct)
egyarr = ne.evaluate('real(-0.5*conj(psi)*funct)')
ekandqlist.append(Vcell * np.sum(egyarr))
tot = tot + Vcell * np.sum(egyarr)
egylist.append(tot)
egyarr = ne.evaluate('real((abs(psi))**2)')
mtotlist.append(Vcell * np.sum(egyarr))
def stripe(signal, spectheight, sigmas, sampdist, eval_range):
"""
Populate an array with time-shifted and windowed versions of an audio
signal. This serves as a precursor for FFT calculation. The first
spectrogram time frame coincides with the first sample in the signal.
Out-of-bounds array entries are assumed as zero.
Parameters
----------
signal : array_like
Audio signal
spectheight : int
Height of the linear-frequency spectrogram
sigmas : float
Number of standard deviations after which to cut the window
sampdist : int
Time intervals to sample the spectrogram
eval_range : slice
Time range of the spectrogram to be computed
Returns
-------
stripeplot : ndarray
Populated array
"""
signal = np.asarray(signal)
pos = np.arange(0, signal.size, sampdist)[eval_range]
stripeplot = pyfftw.zeros_aligned((spectheight * 2, pos.size),
order='F',
dtype='complex128')
window = gauss(np.arange(-spectheight, spectheight), spectheight / sigmas)
for i in range(pos.size):
p = pos[i]
lo = max(0, spectheight - p)
hi = min(2 * spectheight, signal.size - p + spectheight)
winsig = np.zeros(2 * spectheight)
winsig[lo:hi] = (signal[p - spectheight + lo:p - spectheight + hi] *
window[lo:hi])
stripeplot[:, i] = winsig
return stripeplot
def __init__ (self, nFields, layout, state = "configuration"):
"""Pass an iterable that contains the size of each dimension, e.g. [1024, 768], and the state of either \"configurarion\" (default) or \"fourier\"."""
self._state = self.StringToState(state)
self._nFields = int(nFields);
if ( nFields < 1 ):
raise FFTWrongSize("Need to have at least one field. Your nFields: " + str(nFields));
self._layout = [int(x) for x in layout]
self._layoutHalf = [int(x) // 2 for x in layout]
self.nPoints = 1
for x in self._layout:
self.nPoints *= x
self.normalization = 1. / self.nPoints
self.nD = len(self._layout)
# had to follow this: https://github.com/pyFFTW/pyFFTW/issues/29
# // integer division
lastDimSizeComplex = (self._layout[-1] // 2 + 1)
# [n0, n1, n2]
self._layoutMany = self._layout[:]
self._layoutMany[-1] = lastDimSizeComplex;
self._layoutManyRealPadded = self._layoutMany[:]
self._layoutManyRealPadded[-1] *= 2 # symmetry in last dim: z* = z
self._data = pyfftw.zeros_aligned(self._layoutMany + [nFields], dtype = self.__defaultDTypeComplex )
# https://stackoverflow.com/a/12116854 ellipsis
self._realViewPadded = self._data.view(self.__defaultDTypeReal).reshape(*(self._layoutManyRealPadded + [nFields]))
self._realView = self._realViewPadded[..., :self._layout[-1], :nFields]
self._complexView = self._data.view(self.__defaultDTypeComplex)
fftAxes = [x for x in range(len(self._layout))]
self._planR2C = pyfftw.FFTW(self._realView, self._complexView, axes = fftAxes, direction = "FFTW_FORWARD");
self._planC2R = pyfftw.FFTW(self._complexView, self._realView, axes = fftAxes, direction = "FFTW_BACKWARD");
def test_zeros_aligned(self):
shape = (10, 10)
# Test a few alignments and dtypes
for each in [
(3, "float64"),
(7, "float64"),
(9, "float32"),
(16, "int64"),
(24, "bool"),
(23, "complex64"),
(63, "complex128"),
(64, "int8"),
]:
n = each[0]
a = numpy.zeros(shape, dtype=each[1])
b = zeros_aligned(shape, dtype=each[1], n=n)
self.assertTrue(b.ctypes.data % n == 0)
self.assertTrue(b.dtype == each[1])
self.assertTrue(numpy.array_equal(a, b))
def plan_fft2(dims1):
'''Plan FFT for pyfftw for a 2D image.
Inputs:
dims1 (tuplel of int, shape = (2,)): lateral dimension of the image.
Outputs:
bb(2D numpy.ndarray of float32): array of the real video.
bf(2D numpy.ndarray of complex64): array of the complex spectrum.
fft_object_b(pyfftw.FFTW): Object for forward FFT.
fft_object_c(pyfftw.FFTW): Object for inverse FFT.
'''
(rows1, cols1) = dims1
bb = pyfftw.zeros_aligned((rows1, cols1), dtype='float32', n=8)
# No pre-allocation, because this step is not executed in the initialization stage,
# So the running time counts to the total time.
bf = pyfftw.empty_aligned((rows1, cols1//2+1), dtype='complex64', n=8)
fft_object_b = pyfftw.FFTW(bb, bf, axes=(-2, -1), flags=('FFTW_MEASURE',), direction='FFTW_FORWARD', threads=mp.cpu_count())
fft_object_c = pyfftw.FFTW(bf, bb, axes=(-2, -1), flags=('FFTW_MEASURE',), direction='FFTW_BACKWARD', threads=mp.cpu_count())
return bb, bf, fft_object_b, fft_object_c
def check_array(self, array, shape, dtype, copy=True):
"""
Check that a given array is compatible with the FFTW plans,
in terms of alignment and data type.
If the provided array does not meet any of the checks, a new array
is returned.
"""
if array.shape != shape:
raise ValueError("Invalid data shape: expected %s, got %s" %
(shape, array.shape))
if array.dtype != dtype:
raise ValueError("Invalid data type: expected %s, got %s" %
(dtype, array.dtype))
if self.check_alignment and not (pyfftw.is_byte_aligned(array)):
array2 = pyfftw.zeros_aligned(self.shape, dtype=self.dtype_in)
np.copyto(array2, array)
else:
if copy:
array2 = np.copy(array)
else:
array2 = array
return array2
def DownldData4Pressure(self, dataset_name, time, lx, ly, lz, my_id,
auth_token):
chkSz = 32 #This is the maximum possible. May be increased in future depending on network bandwidth
slabs = lx // chkSz
lJHTDB = libJHTDB(auth_token)
lJHTDB.initialize() #NOTE: datbase returns Pressure as [lz,ly,lx]
for k in range(slabs):
start = np.array([my_id * lx + k * chkSz, 0, 0], dtype=np.int)
width = np.array([chkSz, ly, lz], dtype=np.int)
pSlab = lJHTDB.getRawData(time,
start,
width,
data_set=dataset_name,
getFunction='Pressure')
if (k == 0):
pAll = pSlab[:, :, :]
else:
pAll = np.concatenate((pAll, pSlab[:, :, :]),
axis=2) #axis=2=> the index of lx
lJHTDB.finalize()
p = ft.zeros_aligned((lx, ly, lz), dtype='float32')
p[:, :, :] = np.transpose(pAll)
return p
def __init__(self,nx,ny,nz,Nx,Ny,Nz,comm,nproc,rank,fft,dx,dy,dz): self.nx = nx self.ny = ny self.nz = nz self.Nx = Nx self.Ny = Ny self.Nz = Nz self.nproc = nproc self.rank = rank self.comm = comm self.ner = int(nx*np.sqrt(3)) self.rbins = np.linspace(-0.5*dx,2*np.pi*np.sqrt(3)+0.5*dx,ner+1) self.fft = fft self.X = np.zeros((nx,ny,nz), dtype='float32') self.Y = np.zeros((nx,ny,nz), dtype='float32') self.Z = np.zeros((nx,ny,nz), dtype='float32') self.r2 = np.zeros((nx,ny,nz), dtype='float32') self.r2rt = np.zeros((nx,ny,nz), dtype='float32') self.chi = ft.zeros_aligned((nx,ny,nz), dtype='float32') self.chi2 = ft.zeros_aligned((nx,ny,nz), dtype='float32') self.cchi = ft.zeros_aligned((nx,ny,1+(nz//2)), dtype='complex64') self.cchi2 = ft.zeros_aligned((nx,ny,1+(nz//2)), dtype='complex64') self.corr = ft.zeros_aligned((nx,ny,nz),dtype='float32') self.iCorr = ft.zeros_aligned((nx,ny,nz),dtype='float32') self.corrSum = np.zeros(ner,dtype='float32') self.corrF = np.zeros(ner,dtype='float32') self.r2Sum = np.zeros(ner,dtype='float32') self.r2F = np.zeros(ner,dtype='float32') self.corrApend = np.zeros(ner,dtype='float32') self.r2Apend = np.zeros(ner,dtype='float32') self.corrLoc = 0.0 self.r2Loc = 0.0 return
def _ApplyFilter(x, h, *, axes=(-2, -1)):
"""
Applies the filter `h` to the object `x` using the FFTW as follows:
`Y = H . X`
"""
## Initialize ##
# Get the complex shape right #
cs_x = np.array(np.shape(x))
cs_h = np.array(np.shape(h))
cs_x[-1] = cs_x[-1] // 2 + 1
cs_h[-1] = cs_h[-1] // 2 + 1
# Create the FFTW objects #
x_ = fftw.zeros_aligned(np.shape(x), dtype='float32')
h_ = fftw.zeros_aligned(np.shape(h), dtype='float32')
y_ = fftw.zeros_aligned(np.shape(x), dtype='float32')
X = fftw.zeros_aligned(cs_x, dtype='complex64')
H = fftw.zeros_aligned(cs_h, dtype='complex64')
Y = fftw.zeros_aligned(cs_x, dtype='complex64')
FT_X = fftw.FFTW(x_, X, axes=axes)
FT_H = fftw.FFTW(h_, H, axes=axes)
RT_y = fftw.FFTW(Y, y_, axes=axes, direction='FFTW_BACKWARD')
## Filter ##
x_[...] = x
h_[...] = h
FT_X()
FT_H()
Y[...] = X * H
RT_y()
## Output ##
return npf.ifftshift(np.real_if_close(y_), axes=axes)
def evolve(central_mass, num_threads, length, length_units, resol, duration,
duration_units, step_factor, save_number, save_options, save_path,
npz, npy, hdf5, s_mass_unit, s_position_unit, s_velocity_unit,
solitons, start_time):
print('Initialising...')
##########################################################################################
#SET INITIAL CONDITIONS
if (length_units == ''):
gridlength = length
else:
gridlength = convert(length, length_units, 'l')
if (duration_units == ''):
t = duration
else:
t = convert(duration, duration_units, 't')
if (duration_units == ''):
t0 = start_time
else:
t0 = convert(start_time, duration_units, 't')
if (s_mass_unit == ''):
cmass = central_mass
else:
cmass = convert(central_mass, s_mass_unit, 'm')
Vcell = (gridlength / float(resol))**3
ne.set_num_threads(num_threads)
initsoliton_jit = numba.jit(initsoliton)
##########################################################################################
# CREATE THE TIMESTAMPED SAVE DIRECTORY AND CONFIG.TXT FILE
save_path = os.path.expanduser(save_path)
tm = time.localtime()
talt = ['0', '0', '0']
for i in range(3, 6):
if tm[i] in range(0, 10):
talt[i - 3] = '{}{}'.format('0', tm[i])
else:
talt[i - 3] = tm[i]
timestamp = '{}{}{}{}{}{}{}{}{}{}{}{}{}'.format(tm[0], '.', tm[1], '.',
tm[2], '_', talt[0], ':',
talt[1], ':', talt[2], '_',
resol)
file = open('{}{}{}'.format('./', save_path, '/timestamp.txt'), "w+")
file.write(timestamp)
os.makedirs('{}{}{}{}'.format('./', save_path, '/', timestamp))
file = open(
'{}{}{}{}{}'.format('./', save_path, '/', timestamp, '/config.txt'),
"w+")
file.write(('{}{}'.format('resol = ', resol)))
file.write('\n')
file.write(('{}{}'.format('axion_mass (kg) = ', axion_mass)))
file.write('\n')
file.write(('{}{}'.format('length (code units) = ', gridlength)))
file.write('\n')
file.write(('{}{}'.format('duration (code units) = ', t)))
file.write('\n')
file.write(('{}{}'.format('start_time (code units) = ', t0)))
file.write('\n')
file.write(('{}{}'.format('step_factor = ', step_factor)))
file.write('\n')
file.write(('{}{}'.format('central_mass (code units) = ', cmass)))
file.write('\n\n')
file.write(
('{}'.format('solitons ([mass, [x, y, z], [vx, vy, vz], phase]): \n')))
for s in range(len(solitons)):
file.write(('{}{}{}{}{}'.format('soliton', s, ' = ', solitons[s],
'\n')))
file.write(
('{}{}{}{}{}{}'.format('\ns_mass_unit = ', s_mass_unit,
', s_position_unit = ', s_position_unit,
', s_velocity_unit = ', s_velocity_unit)))
file.write(
'\n\nNote: If the above units are blank, this means that the soliton parameters were specified in code units'
)
file.close()
loc = save_path + '/' + timestamp
##########################################################################################
# SET UP THE REAL SPACE COORDINATES OF THE GRID
gridvec = np.linspace(-gridlength / 2.0 + gridlength / float(2 * resol),
gridlength / 2.0 - gridlength / float(2 * resol),
resol)
xarray, yarray, zarray = np.meshgrid(
gridvec,
gridvec,
gridvec,
sparse=True,
indexing='ij',
)
distarray = ne.evaluate(
"(xarray**2+yarray**2+zarray**2)**0.5") # Radial coordinates
##########################################################################################
# SET UP K-SPACE COORDINATES FOR COMPLEX DFT (NOT RHO DFT)
kvec = 2 * np.pi * np.fft.fftfreq(resol, gridlength / float(resol))
kxarray, kyarray, kzarray = np.meshgrid(
kvec,
kvec,
kvec,
sparse=True,
indexing='ij',
)
karray2 = ne.evaluate("kxarray**2+kyarray**2+kzarray**2")
##########################################################################################
# INITIALISE SOLITONS WITH SPECIFIED MASS, POSITION, VELOCITY, PHASE
f = np.load('./Soliton Profile Files/initial_f.npy')
delta_x = 0.00001 # Needs to match resolution of soliton profile array file. Default = 0.00001
warn = 0
psi = pyfftw.zeros_aligned((resol, resol, resol), dtype='complex128')
funct = pyfftw.zeros_aligned((resol, resol, resol), dtype='complex128')
for k in range(len(solitons)):
if (k != 0):
if (not overlap_check(solitons[k], solitons[:k])):
warn = 1
else:
warn = 0
for s in solitons:
mass = convert(s[0], s_mass_unit, 'm')
position = convert(np.array(s[1]), s_position_unit, 'l')
velocity = convert(np.array(s[2]), s_velocity_unit, 'v')
# Note that alpha and beta parameters are computed when the initial_f.npy soliton profile file is generated.
alpha = (mass / 3.883)**2
beta = 2.454
phase = s[3]
funct = initsoliton_jit(funct, xarray, yarray, zarray, position, alpha,
f, delta_x)
####### Impart velocity to solitons in Galilean invariant way
velx = velocity[0]
vely = velocity[1]
velz = velocity[2]
funct = ne.evaluate(
"exp(1j*(alpha*beta*t0 + velx*xarray + vely*yarray + velz*zarray -0.5*(velx*velx+vely*vely+velz*velz)*t0 + phase))*funct"
)
psi = ne.evaluate("psi + funct")
rho = ne.evaluate("real(abs(psi)**2)")
fft_psi = pyfftw.builders.fftn(psi, axes=(0, 1, 2), threads=num_threads)
ifft_funct = pyfftw.builders.ifftn(funct,
axes=(0, 1, 2),
threads=num_threads)
##########################################################################################
# COMPUTE SIZE OF TIMESTEP (CAN BE INCREASED WITH step_factor)
delta_t = (gridlength / float(resol))**2 / np.pi
min_num_steps = t / delta_t
min_num_steps_int = int(min_num_steps + 1)
min_num_steps_int = int(min_num_steps_int / step_factor)
if save_number >= min_num_steps_int:
actual_num_steps = save_number
its_per_save = 1
else:
rem = min_num_steps_int % save_number
actual_num_steps = min_num_steps_int + save_number - rem
its_per_save = actual_num_steps / save_number
h = t / float(actual_num_steps)
##########################################################################################
# SETUP K-SPACE FOR RHO (REAL)
rkvec = 2 * np.pi * np.fft.fftfreq(resol, gridlength / float(resol))
krealvec = 2 * np.pi * np.fft.rfftfreq(resol, gridlength / float(resol))
rkxarray, rkyarray, rkzarray = np.meshgrid(rkvec,
rkvec,
krealvec,
sparse=True,
indexing='ij')
rkarray2 = ne.evaluate("rkxarray**2+rkyarray**2+rkzarray**2")
rfft_rho = pyfftw.builders.rfftn(rho, axes=(0, 1, 2), threads=num_threads)
phik = rfft_rho(rho) # not actually phik but phik is defined in next line
phik = ne.evaluate("-4*3.141593*phik/rkarray2")
phik[0, 0, 0] = 0
irfft_phi = pyfftw.builders.irfftn(phik,
axes=(0, 1, 2),
threads=num_threads)
##########################################################################################
# COMPUTE INTIAL VALUE OF POTENTIAL
phisp = pyfftw.zeros_aligned((resol, resol, resol), dtype='float64')
phisp = irfft_phi(phik)
phisp = ne.evaluate("phisp-(cmass)/distarray")
##########################################################################################
# PRE-LOOP ENERGY CALCULATION
if (save_options[3]):
egylist = []
egpcmlist = []
egpsilist = []
ekandqlist = []
mtotlist = []
calculate_energies(
save_options,
resol,
psi,
cmass,
distarray,
Vcell,
phisp,
karray2,
funct,
fft_psi,
ifft_funct,
egpcmlist,
egpsilist,
ekandqlist,
egylist,
mtotlist,
)
##########################################################################################
# PRE-LOOP SAVE I.E. INITIAL CONFIG
save_grid(
rho,
psi,
resol,
save_options,
npy,
npz,
hdf5,
loc,
-1,
1,
)
##########################################################################################
# LOOP NOW BEGINS
halfstepornot = 1 # 1 for a half step 0 for a full step
tenth = float(
save_number / 10
) #This parameter is used if energy outputs are saved while code is running.
# See commented section below (line 585)
clear_output()
print("The total number of steps is %.0f" % actual_num_steps)
if warn == 1:
print(
"WARNING: Significant overlap between solitons in initial conditions"
)
print('\n')
tinit = time.time()
for ix in range(actual_num_steps):
if halfstepornot == 1:
psi = ne.evaluate("exp(-1j*0.5*h*phisp)*psi")
halfstepornot = 0
else:
psi = ne.evaluate("exp(-1j*h*phisp)*psi")
funct = fft_psi(psi)
funct = ne.evaluate("funct*exp(-1j*0.5*h*karray2)")
psi = ifft_funct(funct)
rho = ne.evaluate("real(abs(psi)**2)")
phik = rfft_rho(
rho) # not actually phik but phik is defined on next line
phik = ne.evaluate("-4*3.141593*(phik)/rkarray2")
phik[0, 0, 0] = 0
phisp = irfft_phi(phik)
phisp = ne.evaluate("phisp-(cmass)/distarray")
#Next if statement ensures that an extra half step is performed at each save point
if (((ix + 1) % its_per_save) == 0) and halfstepornot == 0:
psi = ne.evaluate("exp(-1j*0.5*h*phisp)*psi")
rho = ne.evaluate("real(abs(psi)**2)")
halfstepornot = 1
#Next block calculates the energies at each save, not at each timestep.
calculate_energies(
save_options,
resol,
psi,
cmass,
distarray,
Vcell,
phisp,
karray2,
funct,
fft_psi,
ifft_funct,
egpcmlist,
egpsilist,
ekandqlist,
egylist,
mtotlist,
)
#Uncomment next section if partially complete energy lists desired as simulation runs.
#In this way, some energy data will be saved even if the simulation is terminated early.
# if (save_options[3]):
# if (ix+1) % tenth == 0:
# label = (ix+1)/tenth
# file_name = "{}{}".format(label,'egy_cumulative.npy')
# np.save(os.path.join(os.path.expanduser(loc), file_name), egylist)
# file_name = "{}{}".format(label,'egpcm_cumulative.npy')
# np.save(os.path.join(os.path.expanduser(loc), file_name), egpcmlist)
# file_name = "{}{}".format(label,'egpsi_cumulative.npy')
# np.save(os.path.join(os.path.expanduser(loc), file_name), egpsilist)
# file_name = "{}{}".format(label,'ekandq_cumulative.npy')
# np.save(os.path.join(os.path.expanduser(loc), file_name), ekandqlist)
################################################################################
# SAVE DESIRED OUTPUTS
if ((ix + 1) % its_per_save) == 0:
save_grid(
rho,
psi,
resol,
save_options,
npy,
npz,
hdf5,
loc,
ix,
its_per_save,
)
################################################################################
# UPDATE INFORMATION FOR PROGRESS BAR
tint = time.time() - tinit
tinit = time.time()
prog_bar(actual_num_steps, ix + 1, tint)
################################################################################
# LOOP ENDS
clear_output()
print('\n')
print("Complete.")
if warn == 1:
print(
"WARNING: Significant overlap between solitons in initial conditions"
)
if (save_options[3]):
file_name = "egylist.npy"
np.save(os.path.join(os.path.expanduser(loc), file_name), egylist)
file_name = "egpcmlist.npy"
np.save(os.path.join(os.path.expanduser(loc), file_name), egpcmlist)
file_name = "egpsilist.npy"
np.save(os.path.join(os.path.expanduser(loc), file_name), egpsilist)
file_name = "ekandqlist.npy"
np.save(os.path.join(os.path.expanduser(loc), file_name), ekandqlist)
file_name = "masslist.npy"
np.save(os.path.join(os.path.expanduser(loc), file_name), mtotlist)
ny=1024
nz=1024
idp=8
scrout=sys.stdout
sys.stdout.write('MPI id={0:4d} nx={1:6d} ny={2:6d} nz={3:6d}\n'. \
format(my_id,nx,ny,nz))
lx=nx//nproc
ly=ny
lz=nz
lz_half=lz//2
nek=int(math.sqrt(2.0)/3*nx)
## Initialize the velocity field having a size of (lx,ly,lz)
vx=ft.zeros_aligned((lx,ly,lz), dtype='float32')
vy=ft.zeros_aligned((lx,ly,lz), dtype='float32')
vz=ft.zeros_aligned((lx,ly,lz), dtype='float32')
## Populate velocity field from the Database
comm.Barrier(); t1=MPI.Wtime()
SOAPtdb.loadvel(vx,vy,vz,lx,ly,lz,my_id)
comm.Barrier(); t2=MPI.Wtime()
if(my_id==0):
sys.stdout.write('Load velocity field cost: {0:.2f} seconds\n'.format(t2-t1))
## Transform velocity field from Physical to Fourier space
comm.Barrier(); t1=MPI.Wtime()
myFFT3Dfield=FFT3Dfield()
cvx=myFFT3Dfield.GetFFT3Dfield(vx,lx,ly,lz,nproc,my_id)
cvy=myFFT3Dfield.GetFFT3Dfield(vy,lx,ly,lz,nproc,my_id)
'''
dir_wisdom = 'wisdom\\'
if not os.path.exists(dir_wisdom):
os.makedirs(dir_wisdom)
Dimens = [(504, 504), (464, 504), (416, 480)] # lateral dimension of the video
for ind_video in range(0, len(Dimens)): #
start = time.time()
rows, cols = Dimens[ind_video]
# lateral dimensions slightly larger than the raw video but faster for FFT
x = cv2.getOptimalDFTSize(rows)
y = cv2.getOptimalDFTSize(cols)
# learn 2D wisdom
start1 = time.time()
bb = pyfftw.zeros_aligned((x, y), dtype='float32',
n=8) # numpy array storing the real-space data
bf = pyfftw.zeros_aligned(
(x, y // 2 + 1), dtype='complex64',
n=8) # numpy array storing the Fourier-space data
fft_object_b = pyfftw.FFTW(bb, bf, axes=(-2,-1), flags=('FFTW_MEASURE',), \
direction='FFTW_FORWARD',threads=mp.cpu_count())
fft_object_c = pyfftw.FFTW(bf, bb, axes=(-2,-1), flags=('FFTW_MEASURE',), \
direction='FFTW_BACKWARD',threads=mp.cpu_count())
end1 = time.time()
# Save the learned 2D wisdom result to "wisdom" folder in three ".txt" files
bb = pyfftw.export_wisdom()
print(bb)
Length_data = str((x, y))
file = open(dir_wisdom + Length_data + "x1.txt", "wb")
file.write(bb[0])
import pyfftw
'''Learn wisdom of a 2D array, used for fast FFT planning'''
Dimens = (120,88) # The lateral dimension of the target 2D array
dir_wisdom = 'wisdom'
if not os.path.exists(dir_wisdom):
os.makedirs(dir_wisdom)
start = time.time()
rows, cols = Dimens
x = cv2.getOptimalDFTSize(rows)
y = cv2.getOptimalDFTSize(cols)
start1 = time.time()
bb = pyfftw.zeros_aligned((x, y), dtype='float32', n=8)
bf = pyfftw.zeros_aligned((x, y//2+1), dtype='complex64', n=8)
fft_object_b = pyfftw.FFTW(bb, bf,axes=(-2,-1),flags=('FFTW_MEASURE',), direction='FFTW_FORWARD',threads=mp.cpu_count())
fft_object_c = pyfftw.FFTW(bf, bb,axes=(-2,-1),flags=('FFTW_MEASURE',), direction='FFTW_BACKWARD',threads=mp.cpu_count())
end1 = time.time()
bb = pyfftw.export_wisdom()
print(bb)
Length_data=str((x, y))
file = open(os.path.join(dir_wisdom, Length_data+"x1.txt"), "wb")
file.write(bb[0])
file.close
file = open(os.path.join(dir_wisdom, Length_data+"x2.txt"), "wb")
file.write(bb[1])
file.close
file = open(os.path.join(dir_wisdom, Length_data+"x3.txt"), "wb")
def mtm_spectrum(data, bandwidth, *, fs=1, min_lambda=0.95, n_tapers=None,
remove_mean=False, nfft=None, n_fft_threads=cpu_count()):
"""
Computes the spectrum using Thomson's multitaper method.
Parameters
----------
data : np.ndarray, with shape (T, )
Input data
bandwidth : float
Bandwidth of taper, in Hz
fs : float, optional
Sampling rate, in Hz.
Default is 1 Hz.
min_lambda : float, optional
Minimum energy concentration that each taper must satisfy.
Default is 0.95.
n_tapers : int, optional
Number of tapers to compute
Default is to use all tapers that satisfied 'min_lambda'.
remove_mean : boolean, optional
Whether to remove the mean of the data before computing the
MTM spectrum.
Default is False.
nfft : int, optional
How many FFT points to use for the spectrum.
Default is the same as the length of the input data.
n_fft_threads : int, optional
Number of threads to use for the FFT.
Default is the number of CPUs (which may be virtual).
Returns
-------
mt_sdf : np.ndarray, with shape (nfft, )
The multitapered power spectral density
freqs : np.ndarray, with shape (nfft, )
The corresponding frequencies for the mt PSD, in Hz.
"""
N = data.shape[0]
tapers, lambdas = get_tapers(
N, bandwidth, fs=fs,
n_tapers=n_tapers,
min_lambda=min_lambda)
n_tapers = tapers.shape[0]
if nfft is None:
nfft = N
if remove_mean:
data = data - data.mean()
if np.isrealobj(data):
M = nfft // 2 + 1
xtd = pyfftw.zeros_aligned(
(n_tapers, nfft),
dtype='float64')
xfd = pyfftw.zeros_aligned(
(n_tapers, M),
dtype='complex128')
fft_sig = pyfftw.FFTW(
xtd, xfd,
axes=(1, ),
direction='FFTW_FORWARD',
flags=['FFTW_ESTIMATE'],
threads=n_fft_threads,
planning_timelimit=0)
xtd[:, :N] = tapers * data
xtd[:, N:] = 0
fft_sig(normalise_idft=True)
#assert np.allclose(xfd, np.fft.rfft(tapers * data, n=nfft))
#xfd = np.fft.rfft(tapers * data, n=nfft)
sdfs = (xfd.real**2 + xfd.imag**2) / fs
if nfft % 2 == 0:
sdfs[:, 1:-1] *= 2
else:
sdfs[:, 1:] *= 2
freqs = np.fft.rfftfreq(nfft, d=1/fs)
else:
# can use an in-place transform here
x = pyfftw.zeros_aligned((n_tapers, nfft), dtype='complex128')
fft_sig = pyfftw.FFTW(
x, x,
axes=(1, ),
direction='FFTW_FORWARD',
flags=['FFTW_ESTIMATE'],
threads=n_fft_threads,
planning_timelimit=0)
x[:, :N] = tapers * data
x[:, N:] = 0
fft_sig(normalise_idft=True)
#assert np.allclose(xfd, np.fft.fft(tapers * data, n=nfft))
sdfs = (x.real**2 + x.imag**2) / fs
freqs = np.fft.fftfreq(nfft, d=1/fs)
mt_sdf = np.mean(sdfs, axis=0)
return mt_sdf, freqs
def mtm_spectrogram(data, bandwidth, *, fs=1, timestamps=None, nperseg=None, noverlap=None,
n_tapers=None, min_lambda=0.95, remove_mean=False, nfft=None,
n_fft_threads=cpu_count()):
"""
Computes the spectrogram using Thomson's multitaper method.
Parameters
----------
data : np.ndarray, with shape (T, )
Input data
bandwidth : float
Bandwidth of taper, in Hz
fs : float, optional
Sampling rate, in Hz.
Default is 1 Hz.
timestamps : np.ndarray, with shape (T, ), optional
Timestamps for the data. If not provided, they will be
inferred using np.arange(len(data)) / fs
nperseg : int, optional
Number of samples to use for each segment/window.
Default is 256.
noverlap : int, optional
Number of points to overlap between segments.
Default is nperseg // 8.
min_lambda : float, optional
Minimum energy concentration that each taper must satisfy.
Default is 0.95.
n_tapers : int, optional
Number of tapers to compute
Default is to use all tapers that satisfied 'min_lambda'.
remove_mean : boolean, optional
Whether to remove the mean of the data before computing the
MTM spectrum.
Default is False.
nfft : int, optional
How many FFT points to use for each segment.
Default is the value of 'nperseg'
n_fft_threads : int, optional
Number of threads to use for the FFT.
Default is the number of CPUs (which may be virtual).
Returns
-------
S : np.ndarray, with shape (n_freqs, n_timepoints)
Multitapered spectrogram (units are power spectral density)
f : np.narray, with shape (n_freqs, )
Spectrogram frequencies
t : np.ndarray, with shape (n_timepoints, )
The midpoints of each segment/window.
"""
N = data.shape[0]
if timestamps is None:
timestamps = np.arange(N) / fs
if timestamps.shape[0] != N:
raise ValueError(
f"Expected timestamps to contain {N} elements but got {timestamps.shape[0]}")
estimated_fs = 1.0/np.median(np.diff(timestamps))
if np.abs((estimated_fs - fs)/fs) > 0.01:
print("Warning: estimated fs and provided fs differ by more than 1%")
if nperseg is None:
nperseg = 256
if noverlap is None:
noverlap = nperseg // 8
if noverlap >= nperseg:
raise ValueError("noverlap must be less than {}".format(nperseg))
if nfft is None:
nfft = nperseg
if nfft < nperseg:
raise ValueError(f"'nfft' must be at least {nperseg}")
if nperseg > N:
raise ValueError(f"'nperseg' cannot be larger than the data size {N}")
if not N > noverlap:
raise ValueError(f"'noverlap' cannot be larger than {N-1}")
if remove_mean:
data = data - data.mean()
tapers, lambdas = get_tapers(
nperseg,
bandwidth,
fs=fs,
n_tapers=n_tapers,
min_lambda=min_lambda)
n_tapers = tapers.shape[0]
step = nperseg - noverlap
shape = data.shape[:-1]+((data.shape[-1]-noverlap)//step, nperseg)
strides = data.strides[:-1]+(step*data.strides[-1], data.strides[-1])
data_strided = as_strided(
data,
shape=shape,
strides=strides,
writeable=False)
n_segments = data_strided.shape[0]
out_timestamps = np.mean(
as_strided(
timestamps,
shape=shape,
strides=strides,
writeable=False),
axis=1)
if np.isrealobj(data):
M = nfft // 2 + 1
xtd = pyfftw.zeros_aligned(
(n_tapers, n_segments, nfft),
dtype='float64')
xfd = pyfftw.zeros_aligned(
(n_tapers, n_segments, M),
dtype='complex128')
fft_sig = pyfftw.FFTW(
xtd, xfd,
axes=(2, ),
direction='FFTW_FORWARD',
flags=['FFTW_ESTIMATE'],
threads=n_fft_threads,
planning_timelimit=0)
# (1, n_segments, nperseg) x (n_tapers, 1, nperseg)
xtd[:, :, :N] = data_strided[None, :, :] * tapers[:, None, :]
xtd[:, :, N:] = 0
fft_sig(normalise_idft=True)
#assert np.allclose(xfd, np.fft.rfft(data_strided[None, :, :] * tapers[:, None, :], n=nfft, axis=-1))
#xfd = np.fft.rfft(tapers * data, n=nfft)
spectrograms = (xfd.real**2 + xfd.imag**2) / fs
if nfft % 2 == 0:
spectrograms[:, :, 1:-1] *= 2
else:
spectrograms[:, :, 1:] *= 2
freqs = np.fft.rfftfreq(nfft, d=1/fs)
else:
# can use an in-place transform here
x = pyfftw.zeros_aligned(
(n_tapers, n_segments, nfft),
dtype='complex128')
fft_sig = pyfftw.FFTW(
x, x,
axes=(2, ),
direction='FFTW_FORWARD',
flags=['FFTW_ESTIMATE'],
threads=n_fft_threads,
planning_timelimit=0 )
# (1, n_segments, nperseg) x (n_tapers, 1, nperseg)
x[:, :, :N] = data_strided[None, :, :] * tapers[:, None, :]
x[:, :, N:] = 0
fft_sig(normalise_idft=True)
#assert np.allclose(
# xfd, np.fft.fft(data_strided[None, :, :] * tapers[:, None, :], n=nfft, axis=-1))
spectrograms = (x.real**2 + x.imag**2) / fs
freqs = np.fft.fftfreq(nfft, d=1/fs)
spectrogram = np.sum(lambdas[:, None, None] * spectrograms, axis=0) / np.sum(lambdas)
assert np.all(np.isfinite(spectrogram))
return spectrogram.T, freqs, out_timestamps
def _allocate_array(shape, dtype, fftw):
if fftw:
return zeros_aligned(shape, dtype=dtype, n=simd_alignment)
else:
return np.zeros(shape, dtype)
def alm2rfft(self, alm):
assert alm.size == self.alm_size, alm.size
ret = pyfftw.zeros_aligned(self.ell_mat.rshape, dtype='complex128')
ret[self._cond()] = alm * self.fac_alm2rfft
return ret
def _allocate(self, shape, dtype):
return pyfftw.zeros_aligned(shape, dtype=dtype)
