def __call__(self, r):
"""Calculate the characteristic function from the transform of the BaseModel."""
# Determine q-values.
# We want very fine r-spacing so we can properly normalize f(r). This
# equates to having a large qmax so that the Fourier transform is
# finely spaced. We also want the calculation to be fast, so we pick
# qmax such that the number of q-points is a power of 2. This allows us
# to use the fft.
#
# The initial dr is somewhat arbitrary, but using dr = 0.01 allows for
# the f(r) calculated from a particle of diameter 50, over r =
# arange(1, 60, 0.1) to agree with the sphericalCF with Rw < 1e-4%.
#
# We also have to make a q-spacing small enough to compute out to at
# least the size of the signal.
dr = min(0.01, r[1] - r[0])
ed = 2 * self._model.calculate_ER()
# Check for nans. If we find any, then return zeros.
if numpy.isnan(ed).any():
y = numpy.zeros_like(r)
return y
rmax = max(ed, 2 * r[-1])
dq = pi / rmax
qmax = pi / dr
numpoints = int(2**(ceil(log2(qmax/dq))))
qmax = dq * numpoints
# Calculate F(q) = q * I(q) from model
q = fftfreq(int(qmax/dq)) * qmax
fq = q * self._model.evalDistribution(q)
# Calculate g(r) and the effective r-points
rp = fftfreq(numpoints) * 2 * pi / dq
# Note sine transform = imaginary part of ifft
gr = ifft(fq).imag
# Calculate full-fr for normalization
assert (rp[0] == 0.0)
frp = numpy.zeros_like(gr)
frp[1:] = gr[1:] / rp[1:]
# Inerpolate onto requested grid, do not use data after jump in rp
assert (numpoints % 2 == 0)
nhalf = numpoints / 2
fr = numpy.interp(r, rp[:nhalf], gr[:nhalf])
vmask = (r != 0)
fr[vmask] /= r[vmask]
# Normalize. We approximate fr[0] by using the fact that f(r) is linear
# at low r. By definition, fr[0] should equal 1.
fr0 = 2*frp[2] - frp[1]
fr /= fr0
# Fix potential divide-by-zero issue, fr is 1 at r == 0
fr[~vmask] = 1
return fr
def gvecs(self):
"""
Array with the reduced coordinates of the G-vectors.
.. note::
These are the G-vectors of the FFT box and should be used
when we compute quantities on the FFT mesh.
These vectors differ from the gvecs stored in |GSphere| that
are k-centered and enclosed by a sphere whose radius is defined by ecut.
"""
gx_list = np.rint(fftfreq(self.nx) * self.nx)
gy_list = np.rint(fftfreq(self.ny) * self.ny)
gz_list = np.rint(fftfreq(self.nz) * self.nz)
#print(gz_list, gy_list, gx_list)
gvecs = np.empty((self.size, 3), dtype=np.int)
idx = -1
for gx in gx_list:
for gy in gy_list:
for gz in gz_list:
idx += 1
gvecs[idx, :] = gx, gy, gz
return gvecs
def compare_interpolated_spectrum():
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111)
out = fft(ifftshift(f_full))
freqs = fftfreq(len(f_full), d=0.01) # spacing, Ang
sfreqs = fftshift(freqs)
taper = gauss_taper(freqs, sigma=0.0496) #Ang, corresponds to 2.89 km/s at 5150A.
tout = out * taper
ax.plot(sfreqs, fftshift(tout))
wl_h, fl_h = np.abs(np.load("PH6.8kms_0.01ang.npy"))
wl_l, fl_l = np.abs(np.load("PH2.5kms.npy"))
#end edges
wl_he = wl_h[200:-200]
fl_he = fl_h[200:-200]
interp = Sinc_w(wl_l, fl_l, a=5, window='kaiser')
fl_hi = interp(wl_he)
d = wl_he[1] - wl_he[0]
out = fft(ifftshift(fl_hi))
freqs = fftfreq(len(out), d=d)
ax.plot(fftshift(freqs), fftshift(out))
plt.show()
def applyPreconditioner(self, arr):
ftMap = fft2(arr)
Ny = (arr.shape)[0]
Nx = (arr.shape)[1]
pixScaleX = (
numpy.abs(self.ra1 - self.ra0)
/ Nx
* numpy.pi
/ 180.0
* numpy.cos(numpy.pi / 180.0 * 0.5 * (self.dec0 + self.dec1))
)
pixScaleY = numpy.abs(self.dec1 - self.dec0) / Ny * numpy.pi / 180.0
lx = 2 * numpy.pi * fftfreq(Nx, d=pixScaleX)
ly = 2 * numpy.pi * fftfreq(Ny, d=pixScaleY)
ix = numpy.mod(numpy.arange(Nx * Ny), Nx)
iy = numpy.arange(Nx * Ny) / Nx
lMap = ftMap * 0.0
lMap[iy, ix] = numpy.sqrt(lx[ix] ** 2 + ly[iy] ** 2)
lOverl0 = lMap / self.l0
Fl = self.A * ((lOverl0) ** self.alpha) / (1.0 + lOverl0 ** self.alpha)
filteredFtMap = ftMap / Fl
filteredFtMap[0, 0] = 0.0 # avoids nan and makes mean 0
arr[:, :] = (numpy.real(ifft2(filteredFtMap)))[:, :]
def B_ell(flat_map, component, bin_size, window=None):
"""
Return the binned beam function of the given flat map component,
using ell bins of bin_size. If a window name is given, multiply
the map component by the window function before taking the 2-D
DFT.
The returned beam function is the absolute value of the DFT, so it
has the same units as the map component. The returned bins array
contains the left edges of the bins corresponding to the returned
data: for all i in range(bins.size),
bins[i] <= data[i] < bins[i+1]
"""
ell_x, ell_y= np.meshgrid(fftshift(fftfreq(flat_map.x.size, flat_map.dx()/(2*np.pi))),
fftshift(fftfreq(flat_map.y.size, flat_map.dy()/(2*np.pi))))
ell_x = ell_x.T
ell_y = ell_y.T
ell_r = np.sqrt(ell_x**2 + ell_y**2)
ell_bins = np.arange(0, np.max(ell_r), bin_size)
beam = flat_map[component]
if window is not None:
beam *= get_window(window, flat_map.y.size)
beam *= get_window(window, flat_map.x.size)[:, np.newaxis]
dft = fftshift(fft2(beam))
# Shift the indices down by 1 so that they correspond to the data
# indices. These indices should always be valid indices for the
# DFT data because r_ell has a lower bound at 0 and max(r_ell)
# will have index ell_bins.size.
bin_indices = np.digitize(ell_r.flatten(), ell_bins) - 1
binned = np.zeros(ell_bins.size) # dtype?
for i in range(binned.size):
binned[i] = np.sqrt(np.mean(abs(dft.flatten()[i==bin_indices])**2))
return ell_bins, binned, dft
def __init__(self, cube, header, phys_units=False):
super(VCS, self).__init__()
self.header = header
self.cube = cube
self.fftcube = None
self.correlated_cube = None
self.ps1D = None
self.phys_units = phys_units
if np.isnan(self.cube).any():
self.cube[np.isnan(self.cube)] = 0
# Feel like this should be more specific
self.good_pixel_count = np.sum(self.cube.max(axis=0) != 0)
else:
self.good_pixel_count = float(self.cube.shape[1] * self.cube.shape[2])
# Lazy check to make sure we have units of km/s
if np.abs(self.header["CDELT3"]) > 1:
self.vel_to_pix = np.abs(self.header["CDELT3"]) / 1000.0
else:
self.vel_to_pix = np.abs(self.header["CDELT3"])
self.vel_channels = np.arange(1, self.cube.shape[0], 1)
if self.phys_units:
self.vel_freqs = np.abs(fftfreq(self.cube.shape[0])) / self.vel_to_pix
else:
self.vel_freqs = np.abs(fftfreq(self.cube.shape[0]))
def gen_wedge_psf(nx, ny, nf, dx, dy, df, z, out, threads=None):
u = fftshift(fftfreq(nx, dx * np.pi / 180))
v = fftshift(fftfreq(ny, dy * np.pi / 180))
e = fftshift(fftfreq(nf, df))
E = np.sqrt(Cosmo.Om0 * (1 + z) ** 3 +
Cosmo.Ok0 * (1 + z) ** 2 + Cosmo.Ode0)
D = Cosmo.comoving_transverse_distance(z).value
H0 = Cosmo.H0.value * 1e3
c = const.c.value
print(E, D, H0)
kx = u * 2 * np.pi / D
ky = v * 2 * np.pi / D
k_perp = np.sqrt(kx ** 2 + ky[np.newaxis, ...].T ** 2)
k_par = e * 2 * np.pi * H0 * f21 * E / (c * (1 + z) ** 2)
arr = np.ones((nf, nx, ny), dtype='complex128')
for i in range(nf):
mask = (k_perp > np.abs(k_par[i]) * c * (1 + z) / (H0 * E * D))
arr[i][mask] = 0
np.save('kx.npy', kx)
np.save('ky.npy', ky)
np.save('kpar.npy', k_par)
np.save('wedge_window.npy', arr.real)
fft_arr = fftshift(fftn(ifftshift(arr))).real
hdu = fits.PrimaryHDU(data=fft_arr)
hdr_dict = dict(cdelt1=dx, cdelt2=dy, cdelt3=df,
crpix1=nx/2, crpix2=ny/2, crpix3=nf/2,
crval1=0, crval2=0, crval3=0,
ctype1='RA---SIN', ctype2='DEC--SIN', ctype3='FREQ',
cunit1='deg', cunit2='deg', cunit3='Hz')
for k, v in hdr_dict.items():
hdu.header[k] = v
hdu.writeto(out, clobber=True)
def get_spectrum_1d(data_reg,x_reg,y_reg):
"""Compute the 1d power spectrum.
"""
# remove the mean and squarize
data_reg-=data_reg.mean()
jpj,jpi = data_reg.shape
msize = min(jpj,jpi)
data_reg = data_reg[:msize-1,:msize-1]
x_reg = x_reg[:msize-1,:msize-1]
y_reg = y_reg[:msize-1,:msize-1]
# wavenumber vector
x1dreg,y1dreg = x_reg[0,:],y_reg[:,0]
Ni,Nj = msize-1,msize-1
dx=npy.int(npy.ceil(x1dreg[1]-x1dreg[0]))
k_max = npy.pi / dx
kx = fft.fftshift(fft.fftfreq(Ni, d=1./(2.*k_max)))
ky = fft.fftshift(fft.fftfreq(Nj, d=1./(2.*k_max)))
kkx, kky = npy.meshgrid( ky, kx )
Kh = npy.sqrt(kkx**2 + kky**2)
Nmin = min(Ni,Nj)
leng = Nmin/2+1
kstep = npy.zeros(leng)
kstep[0] = k_max / Nmin
for ind in range(1, leng):
kstep[ind] = kstep[ind-1] + 2*k_max/Nmin
norm_factor = 1./( (Nj*Ni)**2 )
# tukey windowing = tapered cosine window
cff_tukey = 0.25
yw=npy.linspace(0, 1, Nj)
wdw_j = npy.ones(yw.shape)
xw=npy.linspace(0, 1, Ni)
wdw_i= npy.ones(xw.shape)
first_conditioni = xw<cff_tukey/2
first_conditionj = yw<cff_tukey/2
wdw_i[first_conditioni] = 0.5 * (1 + npy.cos(2*npy.pi/cff_tukey * (xw[first_conditioni] - cff_tukey/2) ))
wdw_j[first_conditionj] = 0.5 * (1 + npy.cos(2*npy.pi/cff_tukey * (yw[first_conditionj] - cff_tukey/2) ))
third_conditioni = xw>=(1 - cff_tukey/2)
third_conditionj = yw>=(1 - cff_tukey/2)
wdw_i[third_conditioni] = 0.5 * (1 + npy.cos(2*npy.pi/cff_tukey * (xw[third_conditioni] - 1 + cff_tukey/2)))
wdw_j[third_conditionj] = 0.5 * (1 + npy.cos(2*npy.pi/cff_tukey * (yw[third_conditionj] - 1 + cff_tukey/2)))
wdw_ii, wdw_jj = npy.meshgrid(wdw_j, wdw_i, sparse=True)
wdw = wdw_ii * wdw_jj
data_reg*=wdw
#2D spectrum
cff = norm_factor
tempconj=fft.fft2(data_reg).conj()
tempamp=cff * npy.real(tempconj*fft.fft2(data_reg))
spec_2d=fft.fftshift(tempamp)
#1D spectrum
leng = len(kstep)
spec_1d = npy.zeros(leng)
krange = Kh <= kstep[0]
spec_1d[0] = spec_2d[krange].sum()
for ind in range(1, leng):
krange = (kstep[ind-1] < Kh) & (Kh <= kstep[ind])
spec_1d[ind] = spec_2d[krange].sum()
spec_1d[0] /= kstep[0]
for ind in range(1, leng):
spec_1d[ind] /= kstep[ind]-kstep[ind-1]
return spec_1d, kstep
def test_definition(self):
x = [0, 1, 2, 3, 4, -4, -3, -2, -1]
assert_array_almost_equal(9*fft.fftfreq(9), x)
assert_array_almost_equal(9*pi*fft.fftfreq(9, pi), x)
x = [0, 1, 2, 3, 4, -5, -4, -3, -2, -1]
assert_array_almost_equal(10*fft.fftfreq(10), x)
assert_array_almost_equal(10*pi*fft.fftfreq(10, pi), x)
def source_terms(self, mara, retphi=False):
from numpy.fft import fftfreq, fftn, ifftn
ng = mara.number_guard_zones()
G = self.G
L = 1.0
Nx, Ny, Nz = mara.fluid.shape
Nx -= 2*ng
Ny -= 2*ng
Nz -= 2*ng
P = mara.fluid.primitive[ng:-ng,ng:-ng,ng:-ng]
rho = P[...,0]
vx = P[...,2]
vy = P[...,3]
vz = P[...,4]
K = [fftfreq(Nx)[:,np.newaxis,np.newaxis] * (2*np.pi*Nx/L),
fftfreq(Ny)[np.newaxis,:,np.newaxis] * (2*np.pi*Ny/L),
fftfreq(Nz)[np.newaxis,np.newaxis,:] * (2*np.pi*Nz/L)]
delsq = -(K[0]**2 + K[1]**2 + K[2]**2)
delsq[0,0,0] = 1.0 # prevent division by 0
rhohat = fftn(rho)
phihat = (4*np.pi*G) * rhohat / delsq
fx = -ifftn(1.j * K[0] * phihat).real
fy = -ifftn(1.j * K[1] * phihat).real
fz = -ifftn(1.j * K[2] * phihat).real
S = np.zeros(mara.fluid.shape + (5,))
S[ng:-ng,ng:-ng,ng:-ng,0] = 0.0
S[ng:-ng,ng:-ng,ng:-ng,1] = rho * (fx*vx + fy*vy + fz*vz)
S[ng:-ng,ng:-ng,ng:-ng,2] = rho * fx
S[ng:-ng,ng:-ng,ng:-ng,3] = rho * fy
S[ng:-ng,ng:-ng,ng:-ng,4] = rho * fz
return (S, ifftn(phihat).real) if retphi else S
def subpixel_shift(img_arr, dx=0.5, inplace=True):
'''
Shifts an image by a fraction of a pixel in a random direction,
with circular boundary conditions.
Arguments:
img_arr (2D ndarray): array to subpixel shift in a random direction
dx (float): distance in pixel to shift by
inplace (bool): if True, will modify the array inplace
'''
f = fft.fft2(img_arr, axes=(0, 1))
g = np.meshgrid(
fft.fftfreq(img_arr.shape[-3]),
fft.fftfreq(img_arr.shape[-2])
)
u = npr.normal(0, 1, size=2)
fk = g[0] * u[0] + g[1] * u[1]
phi = np.exp(2j * np.pi * dx * fk)
out = np.real(fft.ifft2(phi[..., np.newaxis] * f, axes=(0, 1)))
if inplace:
img_arr[:] = out[:]
else:
return out
def _configure(infiles, z, df):
conf.infiles = infiles
img_hdr = fits.getheader(conf.infiles[0])
conf.nx = img_hdr['naxis1']
conf.ny = img_hdr['naxis2']
conf.nf = MWA_FREQ_EOR_ALL_80KHZ.size
conf.dx = np.abs(img_hdr['cdelt1']) * np.pi / 180
conf.dy = np.abs(img_hdr['cdelt2']) * np.pi / 180
conf.df = df
conf.du = 1 / (conf.nx * conf.dx)
conf.dv = 1 / (conf.ny * conf.dy)
conf.deta = 1 / (conf.nf * conf.df)
conf.freq = MWA_FREQ_EOR_ALL_80KHZ
conf.u = fftshift(fftfreq(conf.nx, conf.dx))
conf.v = fftshift(fftfreq(conf.ny, conf.dy))
conf.eta = fftshift(fftfreq(conf.nf, conf.df))
conf.z = z
conf.cosmo_d = Cosmo.comoving_transverse_distance(conf.z).value
conf.cosmo_e = np.sqrt(Cosmo.Om0 * (1 + conf.z) ** 3 + Cosmo.Ok0 *
(1 + conf.z) ** 2 + Cosmo.Ode0)
conf.cosmo_h0 = Cosmo.H0.value * 1e3
conf.cosmo_c = const.si.c.value
conf.kx = conf.u * 2 * np.pi / conf.cosmo_d
conf.ky = conf.v * 2 * np.pi / conf.cosmo_d
conf.dkx = conf.du * 2 * np.pi / conf.cosmo_d
conf.dky = conf.dv * 2 * np.pi / conf.cosmo_d
conf.k_perp = np.sqrt(conf.kx ** 2 + conf.ky[np.newaxis, ...].T ** 2)
conf.k_par = conf.eta * 2 * np.pi * conf.cosmo_h0 * _F21 * \
conf.cosmo_e / (conf.cosmo_c * (1 + conf.z) ** 2)
def high_pass_filter(img, filtersize=10):
"""
A FFT implmentation of high pass filter from pyKLIP.
Args:
img: a 2D image
filtersize: size in Fourier space of the size of the space. In image space, size=img_size/filtersize
Returns:
filtered: the filtered image
"""
if filtersize == 0:
return img
# mask NaNs
nan_index = np.where(np.isnan(img))
img[nan_index] = 0
transform = fft.fft2(img)
# coordinate system in FFT image
u,v = np.meshgrid(fft.fftfreq(transform.shape[1]), fft.fftfreq(transform.shape[0]))
# scale u,v so it has units of pixels in FFT space
rho = np.sqrt((u*transform.shape[1])**2 + (v*transform.shape[0])**2)
# scale rho up so that it has units of pixels in FFT space
# rho *= transform.shape[0]
# create the filter
filt = 1. - np.exp(-(rho**2/filtersize**2))
filtered = np.real(fft.ifft2(transform*filt))
# restore NaNs
filtered[nan_index] = np.nan
img[nan_index] = np.nan
return filtered
def fftFromLiteMap(liteMap,applySlepianTaper = False,nresForSlepian=3.0):
"""
@brief Creates an fft2D object out of a liteMap
@param liteMap The map whose fft is being taken
@param applySlepianTaper If True applies the lowest order taper (to minimize edge-leakage)
@param nresForSlepian If above is True, specifies the resolution of the taeper to use.
"""
ft = fft2D()
ft.Nx = liteMap.Nx
ft.Ny = liteMap.Ny
flTrace.issue("flipper.fftTools",1, "Taking FFT of map with (Ny,Nx)= (%f,%f)"%(ft.Ny,ft.Nx))
ft.pixScaleX = liteMap.pixScaleX
ft.pixScaleY = liteMap.pixScaleY
lx = 2*np.pi * fftfreq( ft.Nx, d = ft.pixScaleX )
ly = 2*np.pi * fftfreq( ft.Ny, d = ft.pixScaleY )
ix = np.mod(np.arange(ft.Nx*ft.Ny),ft.Nx)
# This was dividing ints, so change below needed to maintain same behaviour in python3
iy = np.array(np.arange(ft.Nx*ft.Ny)/ft.Nx, dtype = int)
modLMap = np.zeros([ft.Ny,ft.Nx])
modLMap[iy,ix] = np.sqrt(lx[ix]**2 + ly[iy]**2)
ft.modLMap = modLMap
ft.lx = lx
ft.ly = ly
ft.ix = ix
ft.iy = iy
ft.thetaMap = np.zeros([ft.Ny,ft.Nx])
ft.thetaMap[iy[:],ix[:]] = np.arctan2(ly[iy[:]],lx[ix[:]])
ft.thetaMap *=180./np.pi
map = liteMap.data.copy()
#map = map0.copy()
#map[:,:] =map0[::-1,:]
taper = map.copy()*0.0 + 1.0
if (applySlepianTaper) :
try:
path_to_taper = os.path.join(__FLIPPER_DIR__, "tapers", 'taper_Ny%d_Nx%d_Nres%3.1f'%(ft.Ny,ft.Nx,nresForSlepian))
f = open(path_to_taper)
taper = pickle.load(f)
f.close()
except:
taper = slepianTaper00(ft.Nx,ft.Ny,nresForSlepian)
path_to_taper = os.path.join(__FLIPPER_DIR__, "tapers", 'taper_Ny%d_Nx%d_Nres%3.1f'%(ft.Ny,ft.Nx,nresForSlepian))
f = open(path_to_taper, mode="w")
pickle.dump(taper,f)
f.close()
ft.kMap = fftfast.fft(map*taper,axes=[-2,-1])
del map, modLMap, lx, ly
return ft
def icwt2d(W, a, dx=0.25, dy=0.25, da=0.25, wavelet=Mexican_hat()):
"""
Inverse bi-dimensional continuous wavelet transform as in Wang and
Lu (2010), equation [5].
PARAMETERS
W (array like):
Wavelet transform, the result of the cwt2d function.
a (array like, optional):
Scale parameter array.
w (class, optional) :
Mother wavelet class. Default is Mexican_hat()
RETURNS
iW (array like) :
Inverse wavelet transform.
EXAMPLE
"""
m0, l0, k0 = W.shape
if m0 != a.size:
raise Warning('Scale parameter array shape does not match wavelet' \
' transform array shape.')
# Calculates the zonal and meridional wave numters.
L, K = 2 ** int(ceil(log2(l0))), 2 ** int(ceil(log2(k0)))
# Calculates the zonal and meridional wave numbers.
l, k = fftfreq(L, dy), fftfreq(K, dx)
# Creates empty inverse wavelet transform array and fills it for every
# discrete scale using the convolution theorem.
iW = zeros((m0, L, K), 'complex')
for i, an in enumerate(a):
psi_ft_bar = an * wavelet.psi_ft(an * k, an * l)
W_ft = fft2(W[i, :, :], s=(L, K))
iW[i, :, :] = ifft2(W_ft * psi_ft_bar, s=(L, K)) * da / an ** 2.
return iW[:, :l0, :k0].real.sum(axis=0) / wavelet.cpsi
def icwt2d(self, da=0.25):
'''
Inverse bi-dimensional continuous wavelet transform as in Wang and
Lu (2010), equation [5].
Parameters
----------
da : float, optional
Spacing in the frequency axis.
'''
if self.Wf is None:
raise TypeError("Run cwt2D before icwt2D")
m0, l0, k0 = self.Wf.shape
if m0 != self.scales.size:
raise Warning('Scale parameter array shape does not match\
wavelet transform array shape.')
# Calculates the zonal and meridional wave numters.
L, K = 2 ** int(np.ceil(np.log2(l0))), 2 ** int(np.ceil(np.log2(k0)))
# Calculates the zonal and meridional wave numbers.
l, k = fftfreq(L, self.dy), fftfreq(K, self.dx)
# Creates empty inverse wavelet transform array and fills it for every
# discrete scale using the convolution theorem.
self.iWf = np.zeros((m0, L, K), 'complex')
for i, an in enumerate(self.scales):
psi_ft_bar = an * self.wavelet.psi_ft(an * k, an * l)
W_ft = fftn(self.Wf[i, :, :], s=(L, K))
self.iWf[i, :, :] = ifftn(W_ft * psi_ft_bar, s=(L, K)) *\
da / an ** 2.
self.iWf = self.iWf[:, :l0, :k0].real.sum(axis=0) / self.wavelet.cpsi
return self
def calc_mtf(self):
'''
mtf, omega_t, omega_f = calc_mtf(self)
'''
return (fft2(self.causal()),
fftfreq(self.n_causal,self.omega_t),
fftfreq(self.n_bands,self.omega_f))
def fftFromLiteMap(liteMap,applySlepianTaper = False,nresForSlepian=3.0,fftType=None):
"""
@brief Creates an fft2D object out of a liteMap
@param liteMap The map whose fft is being taken
@param applySlepianTaper If True applies the lowest order taper (to minimize edge-leakage)
@param nresForSlepian If above is True, specifies the resolution of the taeper to use.
"""
ft = fft2D()
ft.Nx = liteMap.Nx
ft.Ny = liteMap.Ny
flTrace.issue("flipper.fftTools",1, "Taking FFT of map with (Ny,Nx)= (%f,%f)"%(ft.Ny,ft.Nx))
ft.pixScaleX = liteMap.pixScaleX
ft.pixScaleY = liteMap.pixScaleY
lx = 2*numpy.pi * fftfreq( ft.Nx, d = ft.pixScaleX )
ly = 2*numpy.pi * fftfreq( ft.Ny, d = ft.pixScaleY )
ix = numpy.mod(numpy.arange(ft.Nx*ft.Ny),ft.Nx)
iy = numpy.arange(ft.Nx*ft.Ny)/ft.Nx
modLMap = numpy.zeros([ft.Ny,ft.Nx])
modLMap[iy,ix] = numpy.sqrt(lx[ix]**2 + ly[iy]**2)
ft.modLMap = modLMap
ft.lx = lx
ft.ly = ly
ft.ix = ix
ft.iy = iy
ft.thetaMap = numpy.zeros([ft.Ny,ft.Nx])
ft.thetaMap[iy[:],ix[:]] = numpy.arctan2(ly[iy[:]],lx[ix[:]])
ft.thetaMap *=180./numpy.pi
map = liteMap.data.copy()
#map = map0.copy()
#map[:,:] =map0[::-1,:]
taper = map.copy()*0.0 + 1.0
if (applySlepianTaper) :
try:
f = open(taperDir + os.path.sep + 'taper_Ny%d_Nx%d_Nres%3.1f'%(ft.Ny,ft.Nx,nresForSlepian))
taper = pickle.load(f)
f.close()
except:
taper = slepianTaper00(ft.Nx,ft.Ny,nresForSlepian)
f = open(taperDir + os.path.sep + 'taper_Ny%d_Nx%d_Nres%3.1f'%(ft.Ny,ft.Nx,nresForSlepian),mode="w")
pickle.dump(taper,f)
f.close()
if fftType=='fftw3':
ft.kMap = fft.fft(map*taper,axes=[-2,-1])
else:
ft.kMap = fft2(map*taper)
del map, modLMap, lx, ly
return ft
def _sample_modspec(self,*args,**kwargs):
'''
Return samples of complex modulation spectrum
'''
sliced = self._timeslice(*args,**kwargs)
winlength = sliced.shape[self.time_axis+1]
tf = fftfreq(winlength,1./self.t_rate)
sf = fftfreq(self.n_bands,self.band_spacing)
ft = fft2(sliced)
return ft, tf, sf
def spectral_multiplier(par,dt):
dx=par['dx']
dy=par['dy']
n0=par['n'][0]
n1=par['n'][1]
f0=2.0*numpy.pi*fftfreq(n0,dx)
f1=2.0*numpy.pi*fftfreq(n1,dy)
kx=numpy.outer(f0,numpy.ones(n0))
ky=numpy.outer(numpy.ones(n1),f1)
return numpy.exp(dt*(par['epsilon']-(par['k0']-kx**2-ky**2)**2))
def get_k(self):
"""
"""
azisz = self.get_info('azimuth_size')
azidk = self.get_info('azimuth_dk')
ransz = self.get_info('range_size')
randk = self.get_info('range_dk')
k = (fftshift(fftfreq(azisz))*azisz*azidk,
fftshift(fftfreq(ransz))*ransz*randk)
return k
def __init__ (self, grid):
from numpy import pi, meshgrid
from numpy.fft import fftfreq
kx = 2*pi*fftfreq (grid.nx)/grid.dx
ky = 2*pi*fftfreq (grid.ny)/grid.dy
self.kx, self.ky = meshgrid (kx, ky)
self.phase = self.ky*grid.x
def _dftscore(self, arr):
# Measures the amount of DFT artifacts (in a quite strict manner)
dft = fft.fft2(arr) * self.whatsize
dft /= dft.size
yfreqs = fft.fftfreq(arr.shape[0])[:, np.newaxis]
xfreqs = fft.fftfreq(arr.shape[1])[np.newaxis, :]
weifun = xfreqs ** 2 + yfreqs ** 2
ret = np.abs(dft) * weifun
return ret.sum()
def plot_sinc_windows():
fig, ax = plt.subplots(nrows=2, figsize=(8, 8))
xs = np.linspace(-2, 2., num=200)
xs4 = np.linspace(-5, 5, num=500)
y2s = sinc_w(xs, 'lanczos')
y4s = sinc_w(xs4, 'lanczos', a=5)
yks = sinc_w(xs4, 'kaiser', a=5, alpha=5)
yks2 = sinc_w(xs, 'kaiser', a=2, alpha=5)
ax[0].plot(xs, y2s, "b", label='Lanczos, a=2')
ax[0].plot(xs4, y4s, 'g', label='Lanczos, a=5')
ax[0].plot(xs4, yks, "r", label='Kaiser 5, a=5')
ax[0].plot(xs, yks2, "c", label='Kaiser 5, a=2')
#ax[0].plot(xs,sinc_w(xs, 'hann'),label='Hann')
#ax[0].plot(xs,sinc_w(xs, 'kaiser',alpha=5),label='Kaiser=5')
#ax[0].plot(xs,sinc_w(xs, 'kaiser',alpha=10),label='Kaiser=10')
#xs4 = np.linspace(-4,4,num=100)
#ax[0].plot(xs4,sinc_w(xs4, 'lanczos', a = 4), label='Lanczos,a=4')
ax[0].legend()
ax[0].set_xlabel(r"$\pm a$")
#n=400 #zeropadd FFT
#freqs2 = fftfreq(len(y2s),d=xs[1]-xs[0])
#freqs4 =fftfreq(400,d=xs4[1]-xs4[0])
ysh = ifftshift(y2s)
pady = np.concatenate((ysh[:100], np.zeros((1000,)), ysh[100:]))
freq2 = fftshift(fftfreq(len(pady), d=0.02))
ys4h = ifftshift(y4s)
pad4y = np.concatenate((ys4h[:250], np.zeros((2000,)), ys4h[250:]))
freq4 = fftshift(fftfreq(len(pad4y), d=0.02))
fpady = fft(pady)
fpad4y = fft(pad4y)
ax[1].plot(freq2, 10 * np.log10(np.abs(fftshift(fpady / fpady[0]))))
ax[1].plot(freq4, 10 * np.log10(np.abs(fftshift(fpad4y / fpad4y[0]))))
ysk = ifftshift(yks)
padk = np.concatenate((ysk[:250], np.zeros((2000,)), ysk[250:]))
fpadk = fft(padk)
ax[1].plot(freq4, 10 * np.log10(np.abs(fftshift(fpadk / fpadk[0]))))
ysk2 = ifftshift(yks2)
padk2 = np.concatenate((ysk2[:100], np.zeros((1000,)), ysk2[100:]))
fpadk2 = fft(padk2)
ax[1].plot(freq2, 10 * np.log10(np.abs(fftshift(fpadk2 / fpadk2[0]))))
#ax[1].plot(freqs4, fft(ifftshift(
#ax[1].plot(freqs, get_db_response(xs, 'hann'),label='Hann')
#ax[1].plot(freqs, get_db_response(xs, 'kaiser',alpha=5),label='Kaiser=5')
#ax[1].plot(freqs, get_db_response(xs, 'kaiser',alpha=10),label='Kaiser=10')
#ax[1].legend()
ax[1].set_ylabel("dB")
ax[1].set_xlabel("cycles/a")
plt.show()
def __init__( self, t, x, y, A=[1., 0.85], a=[0.25, 0.85] ):
'''
Initializing generalized thalamo-cortical loop. Full
functionality is only obtained in the subclasses, like DOG,
full_eDOG, etc.
Parameters
----------
t : array
1D Time vector
x : np.array
1D vector for x-axis sampling points
y : np.array
1D vector for y-axis sampling points
Keyword arguments
-----------------
A : sequence (default A = [1., 0.85])
Amplitudes for DOG receptive field for center and surround,
respectively
a : sequence (default a = [0.25, 0.85])
Width parameter for DOG receptive field for center and surround,
respectively
Usage
-----
Look at subclasses for example usage
'''
# Set parameteres as attributes
self.name = 'pyDOG Toolbox'
self.t = t
self.A = A
self.a = a
self.x = x
self.y = y
# Find sampling rates and sampling freqs and
self.nu_xs = 1./(x[1]-x[0])
self.nu_ys = 1./(y[1]-y[0])
self.fs = 1./(t[1]-t[0])
self.f = fft.fftfreq(pl.asarray(t).size, t[1]-t[0])
# fftshift spatial frequency,
self.nu_x = fft.fftfreq(pl.asarray(x).size, x[1]-x[0])
self.nu_y = fft.fftfreq(pl.asarray(y).size, y[1]-y[0])
# Make meshgrids, may come in handy
self._xx, self._yy = pl.meshgrid(self.x, self.y)
self._nu_xx, self._nu_yy = pl.meshgrid(self.nu_x, self.nu_y)
# r is needed for all circular rfs
self.r = pl.sqrt(self._xx**2 + self._yy**2)
self.k = 2 * pl.pi * pl.sqrt(self._nu_xx**2 + self._nu_yy**2)
def define_k(N, D):
"""
Creates the 3d k arrays from N and D
"""
from numpy.fft import fftfreq
from numpy import array, pi, float32
nx, ny, nz = N
dx, dy, dz = D
#Use a matlab like ndgrid to create the wave-vectors with the proper shape to take advantage of pythons vectorization
kx, ky, kz = ndgrid(fftfreq(nx, dx)*2.*pi, fftfreq(ny, dy)*2.*pi, fftfreq(nz, dz)*2.*pi, dtype=float32)
#Output a single vector
return array([kx, ky, kz], dtype = float32)
def __init__(self, size, Lx=1.0, acst=1.0, nu=0.02, t0=1.0, **kwargs):
super(LinearAD, self).__init__()
self.shape = (size,)
self.size = size
self.Lx = Lx
self.acst = acst
self.nu = nu
self.t0 = t0
self.wave_numbers = 2*math.pi/Lx * fft.fftfreq(size) * size
self.laplacian = -(2*math.pi/Lx * fft.fftfreq(size) * size)**2
def get_k1d(self, kmax=None):
"""
"""
azisz = self.get_info('azimuth_size')
azidk = self.get_info('azimuth_dk')
kazi = fftshift(fftfreq(azisz))*azisz*azidk
ransz = self.get_info('range_size')
randk = self.get_info('range_dk')
kran = fftshift(fftfreq(ransz))*ransz*randk
if kmax is not None:
kazi = kazi[abs(kazi) <= kmax]
kran = kran[abs(kran) <= kmax]
k1d = (kazi, kran)
return k1d
def _fft(self, audio_data):
from numpy import fft
amp = fft.rfft(audio_data)
freq = fft.fftfreq(audio_data.shape[-1])[:len(amp)]
return freq, amp.real
def __init__(self, grid, eps=None):
r"""Compute the Fourier transformation of the position space representation
of the kinetic operator :math:`T`.
:param grid: The position space grid :math:`\Gamma` of which we compute
its Fourier transform :math:`\Omega`.
:param eps: The semi-classical scaling parameter :math:`\varepsilon`. (optional)
"""
# Cache the value of eps if given
self._eps = eps
# TODO: Consider class "FourierSpaceGrid"
D = grid.get_dimension()
N = grid.get_number_nodes()
T = grid.get_extensions()
# Fourier domain nodes
prefactors = [t / (2.0 * pi) for t in T]
omega = [fftfreq(n, d=1.0 / n) for n in N]
omega = [o / p for o, p in zip(omega, prefactors)]
# Reshape properly
shape = [-1] + (D - 1) * [1]
omega = [o.reshape(roll(shape, i)) for i, o in enumerate(omega)]
# Compute the dot product of omega with itself
omega_sqr = reduce(add, map(square, omega))
# Fourier space grid axes
self._omega = omega
# Omega dot Omega
self._omega_sqr = omega_sqr
"""
import numpy as np
import matplotlib.pyplot as mp
import numpy.fft as nf
x = np.linspace(-2*np.pi, 2*np.pi, 1000)
# 叠加1000条曲线
y = np.zeros(x.size)
for i in range(1, 1000):
y += 4/(2*i-1)*np.pi * np.sin((2*i-1)*x)
# 对y做傅里叶变换, 绘制频域图像
ffts = nf.fft(y)
# 获取傅里叶变换的频率序列
freqs = nf.fftfreq(x.size, x[1]-x[0])
pows = np.abs(ffts)
mp.figure('FFT', facecolor='lightgray')
mp.subplot(121)
mp.grid(linestyle=':')
mp.plot(x, y,linewidth=2)
mp.subplot(122)
mp.grid(linestyle=':')
mp.plot(freqs[freqs>0], pows[freqs>0],
c='orangered')
# 通过复数数组,经过ifft操作, 得到原函数
y2 = nf.ifft(ffts)
mp.subplot(121)
mp.plot(x, y2, linewidth=7, alpha=0.5)
def read_files(self,
fileName,
fileType,
fMin=None,
fMax=None,
windowFunction=None,
comment='',
filterNegative=False,
extrapolateBand=False,
changeZ=False,
z0=100,
z1=100):
if DEBUG:
print fileType
assert fileType in FILETYPES
if (windowFunction is None):
windowFunction = 'blackman-harris'
self.windowFunction = windowFunction
if (fileType == 'CST_TimeTrace' or fileType == 'Nicolas'):
if fileType == 'CST_TimeTrace':
[inputTrace, outputTrace,
_], self.metaData = readCSTTimeTrace(fileName,
comment=comment)
elif fileType == 'Nicolas':
[inputTrace, outputTrace] = readNicholasTimeTrace(fileName)
if np.mod(len(inputTrace), 2) == 1:
inputTrace = np.vstack([inputTrace[0, :], inputTrace])
inputTrace[0, 0] -= (inputTrace[2, 0] - inputTrace[1, 0])
outputTrace = np.vstack([outputTrace[0, :], outputTrace])
outputTrace[0, 0] -= (outputTrace[2, 0] - outputTrace[1, 0])
#plt.plot(inputTrace[:,0],inputTrace[:,1])
#plt.plot(outputTrace[:,0],outputTrace[:,1])
#plt.show()
self.fAxis = fft.fftshift(
fft.fftfreq(
len(inputTrace) * 2, inputTrace[1, 0] - inputTrace[0, 0]))
self.gainFrequency = fftRatio(outputTrace[:, 1], inputTrace[:, 1])
elif (fileType == 'CST_S11'):
self.fAxis, self.gainFrequency, self.metaData = readCSTS11(
fileName, comment=comment)
elif (fileType == 'VNAHP_S11'):
self.fAxis, self.gainFrequency, self.metaData = readVNAHP(
fileName, comment=comment)
elif (fileType == 'S11_CSV'):
self.fAxis, self.gainFrequency, self.metaData = readCSV(
fileName, comment=comment)
elif (fileType == 'S11_S1P'):
self.fAxis, self.gainFrequency, self.metaData = readS1P(
fileName, comment=comment)
elif (fileType == 'ANRITSU_CSV'):
self.fAxis, self.gainFrequency, self.metaData = readAnritsu(
fileName, comment=comment)
if (fMin is None):
fMin = self.fAxis.min()
if (fMax is None):
fMax = self.fAxis.max()
if (extrapolateBand):
if DEBUG:
print(self.fAxis.min())
print(self.fAxis.max())
if (fMin < self.fAxis.min()):
fitSelection = self.fAxis < self.fAxis.min() + .01
pReal = np.polyfit(self.fAxis[fitSelection],
np.real(self.gainFrequency[fitSelection]),
1)
pImag = np.polyfit(self.fAxis[fitSelection],
np.imag(self.gainFrequency[fitSelection]),
1)
fLow = np.arange(self.fAxis.min(), fMin,
self.fAxis[0] - self.fAxis[1])
self.fAxis = np.hstack([fLow[::-1], self.fAxis])
self.gainFrequency = np.hstack([
pReal[0] * fLow[::-1] + pReal[1] + 1j *
(pImag[0] * fLow[::-1] + pImag[1]), self.gainFrequency
])
#plt.plot(self.fAxis[fitSelection],np.real(self.gainFrequency[fitSelection]),ls='none',marker='o')
#plt.plot(self.fAxis[fitSelection],self.fAxis[fitSelection]*pReal[0]+pReal[1],ls='--',color='r')
#plt.plot(fLow[::-1],fLow[::-1]*pReal[0]+pReal[1],color='k')
#plt.show()
if (fMax > self.fAxis.max()):
fitSelection = self.fAxis > self.fAxis.max() - .01
pReal = np.polyfit(self.fAxis[fitSelection],
np.real(self.gainFrequency[fitSelection]),
1)
pImag = np.polyfit(self.fAxis[fitSelection],
np.imag(self.gainFrequency[fitSelection]),
1)
fHigh = np.arange(self.fAxis.max(), fMax,
self.fAxis[1] - self.fAxis[0])
self.fAxis = np.hstack([self.fAxis, fHigh])
self.gainFrequency = np.hstack([
self.gainFrequency, pReal[0] * fHigh + pReal[1] + 1j *
(pImag[0] * fHigh + pImag[1])
])
selection = np.logical_and(self.fAxis >= fMin, self.fAxis <= fMax)
nf = len(np.where(selection)[0])
if np.mod(nf, 2) == 1:
nf += 1
maxind = np.where(selection)[0].max()
if maxind < len(selection) - 1:
selection[maxind + 1] = True
else:
minind = np.where(selection)[0].min()
if minind > 0:
selection[minind - 1] = True
else:
nf -= 2
selection[maxind] = False
self.fAxis = self.fAxis[selection]
self.gainFrequency = self.gainFrequency[selection]
if (windowFunction == 'blackman-harris'):
wF = signal.blackmanharris(len(self.fAxis))
wF /= np.sqrt(np.mean(wF**2.))
else:
wF = np.ones(len(self.fAxis))
self.tAxis = fft.fftshift(
fft.fftfreq(len(self.fAxis), self.fAxis[1] - self.fAxis[0]))
if (filterNegative):
gainDelay = fft.fftshift(fft.ifft(fft.fftshift(
self.gainFrequency)))
gainDelay[self.tAxis < 0.] = 0.
self.gainFrequency = fft.fftshift(fft.fft(fft.fftshift(gainDelay)))
self.gainDelay = fft.fftshift(
fft.ifft(fft.fftshift(self.gainFrequency * wF)))
if changeZ:
self.change_impedance(z0, z1)
# -*- coding: utf-8 -*-
from __future__ import unicode_literals
import numpy as np
import numpy.fft as nf
import matplotlib.pyplot as mp
times = np.linspace(0, 2 * np.pi, 201)
sigs1 = 4 / (1 * np.pi) * np.sin(1 * times)
sigs2 = 4 / (3 * np.pi) * np.sin(3 * times)
sigs3 = 4 / (5 * np.pi) * np.sin(5 * times)
sigs4 = 4 / (7 * np.pi) * np.sin(7 * times)
sigs5 = 4 / (9 * np.pi) * np.sin(9 * times)
sigs6 = sigs1 + sigs2 + sigs3 + sigs4 + sigs5
freqs = nf.fftfreq(times.size, times[1] - times[0])
ffts = nf.fft(sigs6)
pows = np.abs(ffts)
sigs7 = nf.ifft(ffts).real
mp.figure('FFT', facecolor='lightgray')
mp.subplot(121)
mp.title('Time Domain', fontsize=16)
mp.xlabel('Time', fontsize=12)
mp.ylabel('Signal', fontsize=12)
mp.tick_params(labelsize=10)
mp.grid(linestyle=':')
mp.plot(times, sigs1, label='{:.4f}'.format(1 / (2 * np.pi)))
mp.plot(times, sigs2, label='{:.4f}'.format(3 / (2 * np.pi)))
mp.plot(times, sigs3, label='{:.4f}'.format(5 / (2 * np.pi)))
mp.plot(times, sigs4, label='{:.4f}'.format(7 / (2 * np.pi)))
mp.plot(times, sigs5, label='{:.4f}'.format(9 / (2 * np.pi)))
mp.plot(times, sigs6, label='{:.4f}'.format(1 / (2 * np.pi)))
mp.plot(times,
sigs7,
# ================================================================================================================
# ================================================================================================================
# Compute the Earth's reference frame x, y and z displacements
# True Earth's Heave reference frame displacement. Source: mk4, pg. (28) 7.3 Buoy axes and references
# ================================================================================================================
Epx[i] = Evx[i] * dt[i] # True Earth's X Position in m.
Epy[i] = Evy[i] * dt[i] # True Earth's Y Position in m.
Epz[i] = Evz[i] * dt[i] # True Earth's Heave Position in m.
# ================================================================================================================
# ================================================================================================================
# Fast Fourier Transform in Earth Reference frame and mask.
# Source: Teng 2002 - NDBC Buoys, Pg. 518
# ================================================================================================================
freqs = fftfreq(n, dt[i]) # Set up the FFT Nyquist frequency.
mask = freqs > 0.0 # Mask the freqs for positive values.
# ================================================================================================================
# Calculate the FFT for the Heave and separate the complex conjugate.
# Source: https://www.ndbc.noaa.gov/wavecalc.shtml
# ================================================================================================================
fft_Epz = fft(Epz, n) # FFT for Heave in m/s^2.
fft_Ep = (fft_Epz.real, fft_Epz.imag
) # Separate the real and imaginary numbers.
Epr = fft_Epz.real # Real numbers.
Epi = fft_Epz.imag # imaginary numbers.
# ================================================================================================================
# FFT for the Zenith Deck North-South and East-West Slopes. Source: mk4 pg. 30-31.
# ================================================================================================================
times = []
for i in range(reader.GetTimeSets().GetNumberOfItems()):
array = reader.GetTimeSets().GetItem(i)
for j in range(array.GetNumberOfTuples()):
times.append(array.GetComponent(j, 0))
eigvalues = []
timecoefficients = []
del merger, ds, reader
snaps = Parallel(n_jobs=6, max_nbytes=1e9, verbose=30)(delayed(compute_snapshot)(files[0], time, sys.argv[1:]) for time in times[1:])
#pdb.set_trace()
N = len(times)-1
np.savez('cache_snapshots.npz', snaps=snaps)
fft_values = np.empty((snaps[0].shape[0], N//2))
xf = fftfreq(N, 0.001)
snapshots = np.empty((snaps[0].shape[0], N))
for i, snap in enumerate(snaps):
snapshots[:,i] = snap
ffts = Parallel(n_jobs=6, max_nbytes=1e9, verbose=30, prefer='threads')(delayed(fft)(snapshots[i,:]) for i in range(snapshots.shape[0]))
for i, snap in enumerate(ffts):
fft_values[i, :] = 2.0 / N * np.abs(snap[:N//2])
np.savez(cacheFile, fft_values=fft_values, xf=xf)
else:
data = np.load(cacheFile)
fft_values = data['fft_values']
xf = data['xf']
files = glob.glob(inputFolder + '/*.case')
files.sort()
import numpy as np
from numpy.fft import fft, fftshift, fftfreq
import matplotlib.pyplot as plt
Lx = 2.0
N = 2048
x = np.linspace(-Lx/2, Lx/2, N, endpoint=True)
y = np.ones_like(x)
eps = 0.10 # arbitrary aperture
mask_y = np.abs(x) <= eps
# Spatial frequency sampling
diam = 2*eps # Aperture diameter we are using
N_diam = Lx/diam # How many Diameters is our window [-1, 1]
# Very useful the fftfreq. It calculates the frequencies array for
# a given number of samples N, which span a window of size N_diam
freq = fftshift(fftfreq(n=N, d=N_diam/N))
# Useless sanity checks
plt.figure()
plt.plot(x, y * mask_y)
plt.show()
masked = y * mask_y
pup0 = masked * np.exp(1j * masked)
f = fftshift(fft(pup0, norm='ortho'))
psf0 = (np.abs(f))**2
PEAK = np.max(psf0)
psf0 /= PEAK
para = np.random.rand(3,10)
A = para[0]
w = para[1]
b = para[2]
y = np.zeros(500)
for i,_ in enumerate(y):
ans = 0
for j in range(len(A)):
ans += A[j] * math.sin(w[j]*x[i]+b[j])
y[i] = ans + random.randint(int(-0.05*ans),int(0.05*ans))
#进行FFT分解，尝试寻找dominating frequency
from numpy import fft
data = np.array(y,'f')
yf = fft.fft(data)
f = fft.fftfreq(data.shape[0])
xf = np.arange(len(yf))
plt.plot(f,yf,color='red')
plt.show()
#Keras LSTM
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import LSTM
import time
split_factor = 0.8
split_num = int(len(data_list[0])*split_factor)
cols = train.columns[1:-1]
def cwt(signal, dt, dj=1. / 12, s0=-1, J=-1, wavelet=Morlet()):
"""
Continuous wavelet transform of the signal at specified scales.
Parameters
----------
signal : numpy.ndarray, list
Input signal array
dt : float
Sample spacing.
dj : float, optional
Spacing between discrete scales. Default value is 0.25.
Smaller values will result in better scale resolution, but
slower calculation and plot.
s0 : float, optional
Smallest scale of the wavelet. Default value is 2*dt.
J : float, optional
Number of scales less one. Scales range from s0 up to
s0 * 2**(J * dj), which gives a total of (J + 1) scales.
Default is J = (log2(N*dt/so))/dj.
wavelet : instance of a wavelet class, optional
Mother wavelet class. Default is Morlet wavelet.
Returns
-------
W : numpy.ndarray
Wavelet transform according to the selected mother wavelet.
Has (J+1) x N dimensions.
sj : numpy.ndarray
Vector of scale indices given by sj = s0 * 2**(j * dj),
j={0, 1, ..., J}.
freqs : array like
Vector of Fourier frequencies (in 1 / time units) that
corresponds to the wavelet scales.
coi : numpy.ndarray
Returns the cone of influence, which is a vector of N
points containing the maximum Fourier period of useful
information at that particular time. Periods greater than
those are subject to edge effects.
fft : numpy.ndarray
Normalized fast Fourier transform of the input signal.
fft_freqs : numpy.ndarray
Fourier frequencies (in 1/time units) for the calculated
FFT spectrum.
Example
-------
mother = wavelet.Morlet(6.)
wave, scales, freqs, coi, fft, fftfreqs = wavelet.cwt(var,
0.25, 0.25, 0.5, 28, mother)
"""
n0 = len(signal) # Original signal length.
if s0 == -1: s0 = 2 * dt / wavelet.flambda() # Smallest resolvable scale:
if J == -1: J = int(np.log2(n0 * dt / s0) / dj) # Number of scales
N = 2**(int(np.log2(n0)) + 1) # Next higher power of 2.
signal_ft = fft.fft(signal, N) # Signal Fourier transform
ftfreqs = 2 * np.pi * fft.fftfreq(N, dt) # Fourier angular frequencies
sj = s0 * 2**(np.arange(0, J + 1) * dj) # The scales
freqs = 1 / (wavelet.flambda() * sj) # As of Mallat 1999
# Creates an empty wavlet transform matrix and fills it for every discrete
# scale using the convolution theorem.
W = np.zeros((len(sj), N), 'complex')
for n, s in enumerate(sj):
psi_ft_bar = ((s * ftfreqs[1] * N)**.5 *
np.conjugate(wavelet.psi_ft(s * ftfreqs)))
W[n, :] = fft.ifft(signal_ft * psi_ft_bar, N)
# Checks for NaN in transform results and removes them from the scales,
# frequencies and wavelet transform.
sel = np.logical_not(np.isnan(W).all(axis=1))
sj = sj[sel]
freqs = freqs[sel]
W = W[sel, :]
# Determines the cone-of-influence. Note that it is returned as a function
# of time in Fourier periods. Uses triangualr Bartlett window with non-zero
# end-points.
coi = (n0 / 2. - abs(np.arange(0, n0) - (n0 - 1) / 2))
coi = wavelet.flambda() * wavelet.coi() * dt * coi
#
return W[:, :n0], sj, freqs, coi, dj, s0, J
通过采样数与采样周期求得傅里叶变换分解所得的曲线的频率序列
2.nf.fft(原函数值序列)
通过原函数值的序列经过傅里叶变换得到一个复数数组，复数的模代表的是振幅，复数的辐角代表初相位
3.nf.ifft()
通过一个复数数组经过逆向傅里叶变换得到一个合成的函数值数组
"""
import numpy.fft as nf
import numpy as np
import matplotlib.pyplot as mp
x = np.linspace(0, 4 * np.pi, num=1000)
# y = np.zeros(1000)
y = 4 * np.pi * np.sin(x)
for i in range(2, 10001):
y += 4 / (2 * i - 1) * np.pi * np.sin((2 * i - 1) * x)
mp.subplot(121)
# mp.plot(x, y, label='fang bo')
req = nf.fft(y)
y_ = nf.ifft(req)
mp.plot(x, y, label='fang bo', color='red')
mp.subplot(122)
freq = nf.fftfreq(y_.size, x[1] - x[0])
power = np.abs(req)
mp.plot(freq, power, label='freq', color='red')
mp.legend()
mp.show()
sample_y = signal(sample_t) + state.standard_normal(sample_size)
sample_d = 4. / (sample_size - 1) # Spacing for linspace array
true_signal = signal(sample_t)
fig1, ax1 = plt.subplots()
ax1.plot(sample_t, sample_y, "k.", label="Noisy signal")
ax1.plot(sample_t, true_signal, "k--", label="True signal")
ax1.set_title("Sample signal with noise")
ax1.set_xlabel("Time")
ax1.set_ylabel("Amplitude")
ax1.legend()
spectrum = fft.fft(sample_y)
freq = fft.fftfreq(sample_size, sample_d)
pos_freq_i = np.arange(1, sample_size // 2, dtype=int)
psd = np.abs(spectrum[pos_freq_i])**2 + np.abs(spectrum[-pos_freq_i])**2
fig2, ax2 = plt.subplots()
ax2.plot(freq[pos_freq_i], psd)
ax2.set_title("PSD of the noisy signal")
ax2.set_xlabel("Frequency")
ax2.set_ylabel("Density")
filtered = pos_freq_i[psd < 1e4]
new_spec = np.zeros_like(spectrum)
new_spec[filtered] = spectrum[filtered]
new_spec[-filtered] = spectrum[-filtered]
def test_filters():
"""Test low-, band-, high-pass, and band-stop filters plus resampling."""
rng = np.random.RandomState(0)
sfreq = 100
sig_len_secs = 15
a = rng.randn(2, sig_len_secs * sfreq)
# let's test our catchers
for fl in ['blah', [0, 1], 1000.5, '10ss', '10']:
pytest.raises((ValueError, TypeError),
filter_data,
a,
sfreq,
4,
8,
None,
fl,
1.0,
1.0,
fir_design='firwin')
for nj in ['blah', 0.5]:
pytest.raises(ValueError,
filter_data,
a,
sfreq,
4,
8,
None,
1000,
1.0,
1.0,
n_jobs=nj,
phase='zero',
fir_design='firwin')
pytest.raises(ValueError,
filter_data,
a,
sfreq,
4,
8,
None,
100,
1.,
1.,
fir_window='foo')
pytest.raises(ValueError,
filter_data,
a,
sfreq,
4,
8,
None,
10,
1.,
1.,
fir_design='firwin') # too short
# > Nyq/2
pytest.raises(ValueError,
filter_data,
a,
sfreq,
4,
sfreq / 2.,
None,
100,
1.0,
1.0,
fir_design='firwin')
pytest.raises(ValueError,
filter_data,
a,
sfreq,
-1,
None,
None,
100,
1.0,
1.0,
fir_design='firwin')
# these should work
create_filter(None, sfreq, None, None)
create_filter(a, sfreq, None, None, fir_design='firwin')
create_filter(a, sfreq, None, None, method='iir')
# check our short-filter warning:
with pytest.warns(RuntimeWarning, match='attenuation'):
# Warning for low attenuation
filter_data(a, sfreq, 1, 8, filter_length=256, fir_design='firwin2')
with pytest.warns(RuntimeWarning, match='Increase filter_length'):
# Warning for too short a filter
filter_data(a, sfreq, 1, 8, filter_length='0.5s', fir_design='firwin2')
# try new default and old default
freqs = fftfreq(a.shape[-1], 1. / sfreq)
A = np.abs(fft(a))
kwargs = dict(fir_design='firwin')
for fl in ['auto', '10s', '5000ms', 1024, 1023]:
bp = filter_data(a, sfreq, 4, 8, None, fl, 1.0, 1.0, **kwargs)
bs = filter_data(a, sfreq, 8 + 1.0, 4 - 1.0, None, fl, 1.0, 1.0,
**kwargs)
lp = filter_data(a,
sfreq,
None,
8,
None,
fl,
10,
1.0,
n_jobs=2,
**kwargs)
hp = filter_data(lp, sfreq, 4, None, None, fl, 1.0, 10, **kwargs)
assert_allclose(hp, bp, rtol=1e-3, atol=2e-3)
assert_allclose(bp + bs, a, rtol=1e-3, atol=1e-3)
# Sanity check ttenuation
mask = (freqs > 5.5) & (freqs < 6.5)
assert_allclose(np.mean(np.abs(fft(bp)[:, mask]) / A[:, mask]),
1.,
atol=0.02)
assert_allclose(np.mean(np.abs(fft(bs)[:, mask]) / A[:, mask]),
0.,
atol=0.2)
# now the minimum-phase versions
bp = filter_data(a,
sfreq,
4,
8,
None,
fl,
1.0,
1.0,
phase='minimum',
**kwargs)
bs = filter_data(a,
sfreq,
8 + 1.0,
4 - 1.0,
None,
fl,
1.0,
1.0,
phase='minimum',
**kwargs)
assert_allclose(np.mean(np.abs(fft(bp)[:, mask]) / A[:, mask]),
1.,
atol=0.11)
assert_allclose(np.mean(np.abs(fft(bs)[:, mask]) / A[:, mask]),
0.,
atol=0.3)
# and since these are low-passed, downsampling/upsampling should be close
n_resamp_ignore = 10
bp_up_dn = resample(resample(bp, 2, 1, n_jobs=2), 1, 2, n_jobs=2)
assert_array_almost_equal(bp[n_resamp_ignore:-n_resamp_ignore],
bp_up_dn[n_resamp_ignore:-n_resamp_ignore], 2)
# note that on systems without CUDA, this line serves as a test for a
# graceful fallback to n_jobs=1
bp_up_dn = resample(resample(bp, 2, 1, n_jobs='cuda'), 1, 2, n_jobs='cuda')
assert_array_almost_equal(bp[n_resamp_ignore:-n_resamp_ignore],
bp_up_dn[n_resamp_ignore:-n_resamp_ignore], 2)
# test to make sure our resamling matches scipy's
bp_up_dn = sp_resample(sp_resample(bp,
2 * bp.shape[-1],
axis=-1,
window='boxcar'),
bp.shape[-1],
window='boxcar',
axis=-1)
assert_array_almost_equal(bp[n_resamp_ignore:-n_resamp_ignore],
bp_up_dn[n_resamp_ignore:-n_resamp_ignore], 2)
# make sure we don't alias
t = np.array(list(range(sfreq * sig_len_secs))) / float(sfreq)
# make sinusoid close to the Nyquist frequency
sig = np.sin(2 * np.pi * sfreq / 2.2 * t)
# signal should disappear with 2x downsampling
sig_gone = resample(sig, 1, 2)[n_resamp_ignore:-n_resamp_ignore]
assert_array_almost_equal(np.zeros_like(sig_gone), sig_gone, 2)
# let's construct some filters
iir_params = dict(ftype='cheby1', gpass=1, gstop=20, output='ba')
iir_params = construct_iir_filter(iir_params, 40, 80, 1000, 'low')
# this should be a third order filter
assert iir_params['a'].size - 1 == 3
assert iir_params['b'].size - 1 == 3
iir_params = dict(ftype='butter', order=4, output='ba')
iir_params = construct_iir_filter(iir_params, 40, None, 1000, 'low')
assert iir_params['a'].size - 1 == 4
assert iir_params['b'].size - 1 == 4
iir_params = dict(ftype='cheby1', gpass=1, gstop=20)
iir_params = construct_iir_filter(iir_params, 40, 80, 1000, 'low')
# this should be a third order filter, which requires 2 SOS ((2, 6))
assert iir_params['sos'].shape == (2, 6)
iir_params = dict(ftype='butter', order=4, output='sos')
iir_params = construct_iir_filter(iir_params, 40, None, 1000, 'low')
assert iir_params['sos'].shape == (2, 6)
# check that picks work for 3d array with one channel and picks=[0]
a = rng.randn(5 * sfreq, 5 * sfreq)
b = a[:, None, :]
a_filt = filter_data(a,
sfreq,
4,
8,
None,
400,
2.0,
2.0,
fir_design='firwin')
b_filt = filter_data(b,
sfreq,
4,
8, [0],
400,
2.0,
2.0,
fir_design='firwin')
assert_array_equal(a_filt[:, None, :], b_filt)
# check for n-dimensional case
a = rng.randn(2, 2, 2, 2)
with pytest.warns(RuntimeWarning, match='longer'):
pytest.raises(ValueError, filter_data, a, sfreq, 4, 8,
np.array([0, 1]), 100, 1.0, 1.0)
# check corner case (#4693)
want_length = int(round(_length_factors['hamming'] * 1000. / 0.5))
want_length += (want_length % 2 == 0)
assert want_length == 6601
h = create_filter(np.empty(10000),
1000.,
l_freq=None,
h_freq=55.,
h_trans_bandwidth=0.5,
method='fir',
phase='zero-double',
fir_design='firwin',
verbose=True)
assert len(h) == 6601
h = create_filter(np.empty(10000),
1000.,
l_freq=None,
h_freq=55.,
h_trans_bandwidth=0.5,
method='fir',
phase='zero',
fir_design='firwin',
filter_length='7s',
verbose=True)
assert len(h) == 7001
h = create_filter(np.empty(10000),
1000.,
l_freq=None,
h_freq=55.,
h_trans_bandwidth=0.5,
method='fir',
phase='zero-double',
fir_design='firwin',
filter_length='7s',
verbose=True)
assert len(h) == 8193 # next power of two
comm = MPI.COMM_WORLD
num_processes = comm.Get_size()
rank = comm.Get_rank()
Np = N / num_processes
X = mgrid[rank * Np:(rank + 1) * Np, :N, :N].astype(float) * 2 * pi / N
U = empty((3, Np, N, N))
U_hat = empty((3, N, Np, N / 2 + 1), dtype=complex)
P = empty((Np, N, N))
P_hat = empty((N, Np, N / 2 + 1), dtype=complex)
U_hat0 = empty((3, N, Np, N / 2 + 1), dtype=complex)
U_hat1 = empty((3, N, Np, N / 2 + 1), dtype=complex)
dU = empty((3, N, Np, N / 2 + 1), dtype=complex)
Uc_hat = empty((N, Np, N / 2 + 1), dtype=complex)
Uc_hatT = empty((Np, N, N / 2 + 1), dtype=complex)
curl = empty((3, Np, N, N))
kx = fftfreq(N, 1. / N)
kz = kx[:(N / 2 + 1)].copy()
kz[-1] *= -1
K = array(meshgrid(kx, kx[rank * Np:(rank + 1) * Np], kz, indexing='ij'),
dtype=int)
K2 = sum(K * K, 0, dtype=int)
K_over_K2 = K.astype(float) / where(K2 == 0, 1, K2).astype(float)
kmax_dealias = 2. / 3. * (N / 2 + 1)
dealias = array((abs(K[0]) < kmax_dealias) * (abs(K[1]) < kmax_dealias) *
(abs(K[2]) < kmax_dealias),
dtype=bool)
a = [1. / 6., 1. / 3., 1. / 3., 1. / 6.]
b = [0.5, 0.5, 1.]
def fftn_mpi(u, fu):
sample_rate, noised_sigs = wf.read(
'C:\\Users\\Administrator\\Desktop\\sucai\\da_data\\noised.wav')
times = np.arange(len(noised_sigs)) / sample_rate
mp.figure("Filter", facecolor='lightgray')
mp.subplot(2, 2, 1)
mp.title("Time Domain", fontsize=16)
mp.ylabel("Signal", fontsize=12)
mp.tick_params(labelsize=8)
mp.grid(linestyle=":")
mp.plot(times[:178], noised_sigs[:178], c='dodgerblue', label='Noised Sigs')
mp.legend()
# 获取音频频域信息，绘制频域：频率/能量图像
freqs = nf.fftfreq(times.size, 1 / sample_rate)
noised_ffts = nf.fft(noised_sigs)
noised_pows = np.abs(noised_ffts)
mp.subplot(222)
mp.title("Frequency Domain", fontsize=16)
mp.ylabel("Power", fontsize=12)
mp.tick_params(labelsize=8)
mp.grid(linestyle=":")
mp.semilogy(freqs[freqs > 0],
noised_pows[freqs > 0],
c='orangered',
label='Noised')
mp.legend()
# 将低频噪声去除后绘制音频频域：频率/能量图
fund_freq = freqs[noised_pows.argmax()]
def frequencies(self):
return fftfreq(self.n_fft_samples, 1.0 / self.sampling_frequency)
timeVector = np.arange(0, 1, 1 / Fs)
Audio, Fs = sf.read('guitarra.wav')
#Audio= (np.sin(2*pi*f*timeVector) + np.sin(2*50*pi*f*timeVector) + np.sin(2*100*pi*f*timeVector))/3
#Audio= np.sin(2*pi*f*timeVector)
abc = ola()
abc.run(Audio, 2)
sd.play(abc.y, Fs)
sd.wait()
#sf.write('speech_dobleDuracion.wav',abc.y,Fs)
n = len(Audio)
timeVector = np.arange(0, n * (1 / Fs), 1 / Fs)
n = len(Audio)
frecVector = fftfreq(n)
espectroVector = fft(Audio)
mask = frecVector > 0
truefft = 2.0 * np.abs(espectroVector / n)
plt.figure(3)
plt.plot(frecVector[mask], truefft[mask], linewidth=4)
n1 = len(abc.y)
frecVector1 = fftfreq(n1)
espectroVector1 = fft(abc.y)
mask1 = frecVector1 > 0
truefft1 = 2.0 * np.abs(espectroVector1 / n)
plt.plot(frecVector1[mask1], truefft1[mask1], linewidth=3)
plt.xlabel('Fs' + '[' + str(Fs) + 'Hz' + ']')
plt.ylabel('Potencia')
z = np.real(sig[:, 0, :])
lim = np.max(np.abs(z))
plt.pcolormesh(expand_axis(t_1),
expand_axis(t_3),
z.T,
cmap=plt.cm.seismic,
vmin=-lim,
vmax=lim)
levels = np.linspace(-lim, lim, 11)
plt.contour(t_1, t_3, z.T, levels=levels, colors='0.8', linewidths=0.5)
outname = fn[:-4] + ".png"
plt.tight_layout()
plt.subplots_adjust(left=0.15, right=0.9, top=0.9, bottom=0.15)
plt.savefig(outname)
f_1 = phz2ev(fftshift(fftfreq(t_1.size, t_1[1] - t_1[0])))
#f_2 = phz2ev(fftshift(fftfreq(t_2.size, t_2[1]-t_2[0])))
f_3 = phz2ev(fftshift(fftfreq(t_3.size, t_3[1] - t_3[0])))
for fresp in ["sr", "sn", "sa"]:
fn = "{}_{}.bin".format(rootname, fresp)
sig = np.memmap(pth.join(fld, fn),
dtype=np.complex128,
shape=(t1_n, t2_n, t3_n),
order="F")
plt.figure(figsize=(4, 4), dpi=300)
z = np.real(sig[:, 0, :])
lim = np.max(np.abs(z))
plt.pcolormesh(expand_axis(f_1),
expand_axis(f_3),
z.T,
xn = x - np.mean(x)
yn = y - np.mean(y)
zn = z - np.mean(z)
# ------------------------------------
# DFT
# ------------------------------------
#calculate time step:
n = len(x)
timeStep = t / n
zf = fft.fft(zn) / n
xf = fft.fft(xn) / n
yf = fft.fft(yn) / n
freq = fft.fftfreq(n, d=timeStep)
# ------------------------------------
# ZAPIS
# ------------------------------------
filename = strftime("%d_%H_%M_%S", gmtime())
np.savetxt('logs/' + filename + '.txt',
(np.real(xf), np.real(yf), np.real(zf), freq),
fmt='%.7e',
newline='\n')
print("Files saved ({}).\n".format(filename))
# ------------------------------------
def simulator(spectrum,
isotopomers,
transitions=[[-0.5, 0.5]],
nt=90,
number_of_sidebands=128):
B0 = spectrum["magnetic_flux_density"]
spin_frequency = spectrum["rotor_frequency"]
rotor_angle = spectrum["rotor_angle"]
frequency_scaling_factor = spectrum["gyromagnetic_ratio"] * B0
number_of_points = spectrum["number_of_points"]
spectral_width = spectrum["spectral_width"]
reference_offset = spectrum["reference_offset"]
frequency = (np.arange(number_of_points) / number_of_points) - 0.5
frequency *= spectral_width
frequency += reference_offset
increment = frequency[1] - frequency[0]
if spin_frequency < 1.0e-3:
spin_frequency = 1.0e9
rotor_angle = 0.0
number_of_sidebands = 1
shift_half_bin = 0.5
# orientations
cos_alpha, cos_beta, orientation_amp = polar_coordinates(nt)
orientation_amp /= np.float64(number_of_sidebands)
n_orientation = cos_beta.size
# sideband freq
vr_freq = fftfreq(number_of_sidebands, d=1.0 / number_of_sidebands)
vr_freq *= spin_frequency / increment
# wigner matrix
wigner_2 = wigner_matrices(2, cos_beta)
# rotor to lab frame transformation
lab_vector_2 = rotation_lab(2, rotor_angle)
# pre phase
pre_phase = pre_phase_components(number_of_sidebands, spin_frequency)
# allocate memory for calculations
R2_out = np.empty((n_orientation, 5), dtype=np.complex128)
spectrum = np.zeros(number_of_points)
shape = (n_orientation, number_of_sidebands)
temp = np.empty(shape, dtype=np.complex128)
sideband_amplitude = np.empty(shape, dtype=np.float64)
local_frequency = np.empty(n_orientation, dtype=np.float64)
freq_offset = np.empty(n_orientation, dtype=np.float64)
offset = np.empty(number_of_sidebands, dtype=np.float64)
# start calculating the spectrum for every site in every isotopomer.
for isotopomer in isotopomers:
sites = isotopomer["sites"]
spec = np.zeros(number_of_points)
for transition in transitions:
for site in sites:
iso = site["isotropic_chemical_shift"]
if iso.unit.physical_type == "dimensionless":
iso = iso.value * frequency_scaling_factor
else:
iso = iso.value
zeta = site["shielding_symmetric"]["anisotropy"]
if zeta.unit.physical_type == "dimensionless":
zeta = zeta.value * frequency_scaling_factor
else:
zeta = zeta.value
eta = site["shielding_symmetric"]["asymmetry"]
# Hailtonian
# nuclear shielding
R0, R2 = NS(iso, zeta, eta)
scale = tf.p(transition[1], transition[0])
R0 *= scale
R2 *= scale
local_frequency_offset = (shift_half_bin +
(R0 - frequency[0]) / increment)
# rotation from PAS to Rotor frame over all orientations
R2_out = rotation(2,
R2,
wigner_matrix=wigner_2,
cos_alpha=cos_alpha)
# rotation from rotor to lab frame over all orientations
R2_out *= lab_vector_2
# calculating side-band amplitudes
temp[:] = fft(exp(dot(R2_out, pre_phase[2:7])), axis=-1)
sideband_amplitude[:] = temp.real**2 + temp.imag**2
sideband_amplitude *= orientation_amp[:, np.newaxis]
# calculating local frequencies
local_frequency[:] = R2_out[:, 2].real / increment
# interpolate in-between the frequencies to generate a smooth spectrum.
offset[:] = vr_freq + local_frequency_offset
# print("before", default_timer() - start0)
for j, shift in enumerate(offset):
if int(shift) >= 0 and int(shift) <= number_of_points:
freq_offset[:] = shift + local_frequency
# This is the slowest part of the code.
averager(spec, freq_offset, nt, sideband_amplitude[:,
j])
# np.vectorize(averager(spec, freq_offset,
# nt, sideband_amplitude[:, j]))
# print("time for computing site", default_timer() - start0)
# average over all spins
spectrum += spec * isotopomer["abundance"]
return frequency, spectrum
f = gauss_1d(x0, a=a, w=px2m(w) / 4, x0=px2m(0))
aperture = rect_1d(x0, w=px2m(width - 10))
# todo
# f2 = np.zeros((x0.size * 2))
# f2[:f.size] = f
# f2[f.size:] = f[::-1]
# f, x = f2, x2
x = x0
r = np.sqrt(x**2 + r0**2)
phase = wave_number * r
field = np.sqrt(f) * np.exp(-1j * phase) # * aperture
nu_x = fftfreq(f.size, d=dx)
kx = 1j * 2 * np.pi * nu_x
exp_term = np.sqrt(1 - (wavelength * nu_x)**2)
h = np.exp(1j * wave_number * mm2m(z) * exp_term)
field_z = ifft(fft(field) * h)
f_z, phase_z = np.abs(field_z)**2, np.angle(field_z)
np_gradient = np.gradient(f_z, dx)
fft_gradient = ifft(fft(f_z) * kx)
fft_gradient_real = real(ifft(fft(f_z) * kx))
error = np.abs(fft_gradient - np_gradient)
# error_real = np.abs(fft_gradient_real-np_gradient)
# error[error < 1e-20] = 1e-15
def __init__(self,
ase_cell: Atoms,
vspin: int,
n1: int,
n2: int,
n3: int,
atomic_density_files: dict = None):
# define cell
self.R = ase_cell.get_cell() * angstrom_to_bohr
self.G = 2 * np.pi * np.linalg.inv(self.R).T
assert np.all(
np.isclose(np.dot(self.R, self.G.T), 2 * np.pi * np.eye(3)))
self.omega = np.linalg.det(self.R)
self.atoms = list(
Atom(sample=self, ase_atom=atom) for atom in ase_cell)
self.natoms = len(self.atoms)
self.species = sorted(set([atom.symbol for atom in self.atoms]))
self.nspecies = len(self.species)
# define fragments and constraints list
self.fragments = []
self.constraints = []
# define vspin and FFT grid
self.vspin = vspin
self.n1, self.n2, self.n3 = n1, n2, n3
self.n = n1 * n2 * n3
# define energies, forces and wavefunction
self.Ed = None
self.Ec = None
self.W = None
self.Fd = None
self.Fc = None
self.Fw = None
self.wfc = None
# define charge density and promolecule charge densities
self.rho_r = None
self.rhopro_tot_r = None
self.rhoatom_g = {}
self.rhoatom_rd = {}
# compute the norm of all G vectors on the [n1, n2, n3] grid
G1, G2, G3 = self.G
# all Gx,Gy,Gz values
G1s = np.outer(G1, fftfreq(n1, d=1. / n1))
G2s = np.outer(G2, fftfreq(n2, d=1. / n2))
G3s = np.outer(G3, fftfreq(n3, d=1. / n3))
# make grid from G1s, G2s,G3s using Python built-in broadcasting
self.Gx_g = (G1s[0, :, np.newaxis, np.newaxis] +
G2s[0, np.newaxis, :, np.newaxis] +
G3s[0, np.newaxis, np.newaxis, :])
self.Gy_g = (G1s[1, :, np.newaxis, np.newaxis] +
G2s[1, np.newaxis, :, np.newaxis] +
G3s[1, np.newaxis, np.newaxis, :])
self.Gz_g = (G1s[2, :, np.newaxis, np.newaxis] +
G2s[2, np.newaxis, :, np.newaxis] +
G3s[2, np.newaxis, np.newaxis, :])
self.G2_g = self.Gx_g**2 + self.Gy_g**2 + self.Gz_g**2
# compute atomic density for all species
if atomic_density_files is not None:
# read atomic density from file
for s in self.species:
rho_r, ase_cell = read_cube_data(atomic_density_files[s])
omega = ase_cell.get_volume() * angstrom_to_bohr**3
assert rho_r.shape == (n1, n2, n3)
# in order to have rhoatom_g on commensurate grid
# as generated by fftfreq, roll around
rho_r1 = np.roll(rho_r, n1 // 2, axis=0)
rho_r2 = np.roll(rho_r1, n2 // 2, axis=1)
rho_r3 = np.roll(rho_r2, n3 // 2, axis=2)
# self.rhoatom_g[s] = omega / self.n * fftn(rho_r3)
self.rhoatom_g[s] = omega / self.n * fftn(rho_r3)
else:
# calculate atomic density from pre-computed spherically-averaged atomic density
# located in atomic/rho
# define mapping from all G vectors to G vectors with unique norm
# i.e., repeated |G| are excluded
# tocheck: is 5A cutoff for atomic densities sufficient
# tocheck: is 0.02 integration step sufficient
# rd_grid size set to [251,]; 5 Ang cutoff with 0.02 Ang step
# Gmapping, rho_g: n1 x n2 x n3
# sinrG: rd_grid x unique |G|
G2_d = np.sort(np.array(list(set(self.G2_g.flatten()))))
self.Gmapping = np.searchsorted(G2_d, self.G2_g)
self.G_d = np.sqrt(G2_d)
self.sinrG = np.sin(np.outer(rd_grid, self.G_d))
# rho(G) = 4 pi int( rho(r) r sinGr / G )
# rho(G=0) = 4 pi int( rho(r) r^2 ) = nel / omega
for s in self.species:
rho_rd = np.loadtxt("{}/{}.spavr".format(rho_path, s),
dtype=float)[:, 1]
rho_rd[rho_rd < 0] = 0
with warnings.catch_warnings():
warnings.simplefilter("ignore")
rho_d = 4 * np.pi * drd * np.einsum(
"r,rg->g", rho_rd * rd_grid, self.sinrG / self.G_d)
rho_g = rho_d[self.Gmapping]
rho_g[0, 0, 0] = 4 * np.pi * drd * np.sum(rho_rd * rd_grid**2)
fac = SG15PP[s]["nel"] / rho_g[0, 0, 0]
rho_rd *= fac
rho_g *= fac # normalize charge density
self.rhoatom_g[s] = rho_g
self.rhoatom_rd[s] = rho_rd
def interpolate_subband(self, nfi, df, f0, full_output=False):
'''
interpolates a sub-band between fMin and fMax by
taking a windowed FFT at a bandwidth twice the
requested bandwidth, extrapolating and interpolating
the delay-transform transform, and FT-ing back.
'''
change_nfi = False
fMin = f0 - nfi / 2 * df
if fMin < self.fAxis.min():
fMin = self.fAxis.min()
change_nfi = True
fMax = f0 + (nfi / 2 - 1) * df
if fMax > self.fAxis.max():
fMax = self.fAxis.max()
change_nfi = True
if change_nfi:
nfi = int(np.round(fMax / df - fMin / df))
f0 = fMin + nfi / 2 * df
fAxis_interp = f0 + np.arange(-nfi / 2, nfi / 2) * df
tAxis_interp = fft.fftshift(fft.fftfreq(nfi, df))
b = fMax - fMin
select_max = np.min([self.fAxis.max(), f0 + b])
select_min = np.max([self.fAxis.min(), f0 - b])
selection = np.logical_and(self.fAxis >= select_min,
self.fAxis <= select_max)
nf = len(self.fAxis[selection])
if np.mod(nf, 2) == 1:
nf += 1
maxind = np.where(selection)[0].max()
if maxind < len(selection) - 1:
selection[maxind + 1] = True
else:
minind = np.where(selection)[0].min()
if minind > 0:
selection[minind - 1] = True
else:
nf -= 2
selection[maxind] = False
sub_band = self.gainFrequency[selection]
sub_fAxis = self.fAxis[selection]
window = signal.blackmanharris(nf)
delay_band = fft.fftshift(fft.ifft(fft.fftshift(sub_band * window)))
sub_tAxis = fft.fftshift(
fft.fftfreq(len(sub_band), self.fAxis[1] - self.fAxis[0]))
maxTime = sub_tAxis.max()
minTimeExt = maxTime * 1. / 3.
maxTimeExt = maxTime * 2. / 3.
ext_select = np.logical_and(sub_tAxis <= maxTimeExt,
sub_tAxis >= minTimeExt)
ext_poly = np.polyfit(sub_tAxis[ext_select],
np.log10(np.abs(delay_band[ext_select])), 1)
interp_func_abs = interp.interp1d(sub_tAxis,
np.log10(np.abs(delay_band)))
interp_func_arg = interp.interp1d(sub_tAxis, np.angle(delay_band))
band_interp = np.zeros(nfi, dtype=complex)
select_interp = np.logical_and(tAxis_interp >= 0,
tAxis_interp < maxTimeExt)
band_interp[select_interp] = 10**(interp_func_abs(
tAxis_interp[select_interp])) * np.exp(
1j * interp_func_arg(tAxis_interp[select_interp]))
select_ext = tAxis_interp >= maxTimeExt
band_interp[select_ext] = 10**(tAxis_interp[select_ext] * ext_poly[0] +
ext_poly[1])
#band_interp[select_ext]=0.
window_interp_func = interp.interp1d(
sub_fAxis, signal.blackmanharris(len(sub_band)))
wFactor = 1. / ((fAxis_interp.max() - fAxis_interp.min()) /
(sub_fAxis.max() - sub_fAxis.min()))
if DEBUG:
print(wFactor)
band_interp_f = fft.fftshift(fft.fft(
fft.fftshift(band_interp))) * wFactor
window_corr = window_interp_func(fAxis_interp)
#band_interp_f/=window_corr
if full_output:
return sub_tAxis, delay_band, sub_fAxis, sub_band, tAxis_interp, band_interp, fAxis_interp, band_interp_f
else:
return fAxis_interp, band_interp_f
def main():
from matplotlib.pyplot import semilogx, plot, show, xlim, ylim, figure, legend, subplot, bar
from numpy.fft import fft, fftfreq, fftshift, ifft
from numpy import log10, linspace, interp, angle, array, concatenate
N = 2048 * 2 * 2
fs = float(SAMPLING_RATE)
Nchannels = 20
low_freq = 20.
impulse = zeros(N)
impulse[N / 2] = 1
f = 1000.
#impulse = sin(2*pi*f*arange(0, N/fs, 1./fs))
#[ERBforward, ERBfeedback] = MakeERBFilters(fs, Nchannels, low_freq)
#y = ERBFilterBank(ERBforward, ERBfeedback, impulse)
BandsPerOctave = 3
Nbands = NOCTAVE * BandsPerOctave
[B, A, fi, fl, fh] = octave_filters(Nbands, BandsPerOctave)
y, zfs = octave_filter_bank(B, A, impulse)
#print "Filter lengths without decimation"
#for b, a in zip(B, A):
# print len(b), len(a)
response = 20. * log10(abs(fft(y)))
freqScale = fftfreq(N, 1. / fs)
figure()
subplot(211)
for i in range(0, response.shape[0]):
semilogx(freqScale[0:N / 2], response[i, 0:N / 2])
xlim(fs / 2000, fs)
ylim(-70, 10)
subplot(212)
m = 0
for f in fi:
p = 10. * log10((y[m]**2).mean())
m += 1
semilogx(f, p, 'ko')
Ndec = 3
fc = 0.5
# other possibilities
#(bdec, adec) = ellip(Ndec, 0.05, 30, fc)
#print bdec
#(bdec, adec) = cheby1(Ndec, 0.05, fc)
#(bdec, adec) = butter(Ndec, fc)
(bdec, adec) = iirdesign(0.48,
0.50,
0.05,
70,
analog=0,
ftype='ellip',
output='ba')
#bdec = firwin(30, fc)
#adec = [1.]
figure()
subplot(211)
response = 20. * log10(abs(fft(impulse)))
plot(fftshift(freqScale), fftshift(response), label="impulse")
y = lfilter(bdec, adec, impulse)
response = 20. * log10(abs(fft(y)))
plot(fftshift(freqScale), fftshift(response), label="lowpass")
ydec = y[::2].repeat(2)
response = 20. * log10(abs(fft(ydec)))
plot(fftshift(freqScale),
fftshift(response),
label="lowpass + dec2 + repeat2")
ydec2 = interp(list(range(0, len(y))), list(range(0, len(y), 2)), y[::2])
response = 20. * log10(abs(fft(ydec2)))
plot(fftshift(freqScale),
fftshift(response),
label="lowpass + dec2 + interp2")
ydec3 = y[::2]
response = 20. * log10(abs(fft(ydec3)))
freqScale2 = fftfreq(N / 2, 2. / fs)
plot(freqScale2, fftshift(response), label="lowpass + dec2")
legend(loc="lower left")
subplot(212)
plot(list(range(0, len(impulse))), impulse, label="impulse")
plot(list(range(0, len(impulse))), y, label="lowpass")
plot(list(range(0, len(impulse))), ydec, label="lowpass + dec2 + repeat2")
plot(list(range(0, len(impulse))), ydec2, label="lowpass + dec2 + interp2")
plot(list(range(0, len(impulse), 2)), ydec3, label="lowpass + dec2")
legend()
[boct, aoct, fi, flow,
fhigh] = octave_filters_oneoctave(Nbands, BandsPerOctave)
y, dec, zfs = octave_filter_bank_decimation(bdec, adec, boct, aoct,
impulse)
#print "Filter lengths with decimation"
#print len(bdec), len(adec)
#for b, a in zip(boct, aoct):
# print len(b), len(a)
figure()
subplot(211)
for yone, d in zip(y, dec):
response = 20. * log10(abs(fft(yone)) * d)
freqScale = fftfreq(N / d, 1. / (fs / d))
semilogx(freqScale[0:N / (2 * d)], response[0:N / (2 * d)])
xlim(fs / 2000, fs)
ylim(-70, 10)
subplot(212)
m = 0
for i in range(0, NOCTAVE):
for f in fi:
p = 10. * log10((y[m]**2).mean())
semilogx(f / dec[m], p, 'ko')
m += 1
[boct, aoct, fi, flow,
fhigh] = octave_filters_oneoctave(Nbands, BandsPerOctave)
y1, dec, zfs = octave_filter_bank_decimation(bdec, adec, boct, aoct,
impulse[0:N / 2])
y2, dec, zfs = octave_filter_bank_decimation(bdec,
adec,
boct,
aoct,
impulse[N / 2:],
zis=zfs)
y = []
for y1one, y2one in zip(y1, y2):
y += [concatenate((y1one, y2one))]
figure()
subplot(211)
for yone, d in zip(y, dec):
response = 20. * log10(abs(fft(yone)) * d)
freqScale = fftfreq(N / d, 1. / (fs / d))
semilogx(freqScale[0:N / (2 * d)], response[0:N / (2 * d)])
xlim(fs / 2000, fs)
ylim(-70, 10)
subplot(212)
m = 0
for i in range(0, NOCTAVE):
for f in fi:
p = 10. * log10((y[m]**2).mean())
semilogx(f / dec[m], p, 'ko')
m += 1
generate_filters_params()
show()
flsig4 = (bpf1(data))
flsig5 = (bpf2(data))
flsig6 = (bpf3(data))
plt.title("segmented signal")
plt.plot(flsig4, color='green')
plt.plot(flsig5, color='blue')
plt.plot(flsig6, color='red')
plt.show()
datafft4 = fft_plot(flsig4, "fft :: song 1 no shift")
datafft5 = fft_plot(flsig5, "fft :: song 2 no shift")
datafft6 = fft_plot(flsig6, "fft :: song 3 no shift")
plt.title("segmented signals :: Fourier Transform")
freq = fft.fftshift(fft.fftfreq(datafft4.shape[-1]))
plt.plot(freq, np.abs(datafft4), color='green')
plt.plot(freq, np.abs(datafft5), color='blue')
plt.plot(freq, np.abs(datafft6), color='red')
plt.show()
flsig4 = lpf(flsig4 * w1)
flsig5 = lpf(flsig5 * w2)
flsig6 = lpf(flsig6 * w3)
plt.title("segmented signal")
plt.plot(flsig4, color='green')
plt.plot(flsig5, color='blue')
plt.plot(flsig6, color='red')
plt.show()
def compute_fta(xcor,
time,
fs,
distance,
alpha,
periods,
vmin=1,
vmax=5,
snr_threshold=4.0):
# functions used within Frequency Time Analysis routine
def compute_instantaneous(xcor, fs):
analytic = signal.hilbert(xcor)
env = abs(analytic)
instantaneous_phase = np.unwrap(np.angle(analytic))
instantaneous_frequency = (np.diff(instantaneous_phase) /
(2.0 * np.pi) * fs)
instantaneous_frequency = np.append(instantaneous_frequency, 0)
# to have the same number of points
return env, instantaneous_frequency
def split_trace(trace, time):
n = int(np.floor((len(time)) / 2)) # index of 0 value on time axis
matrix = np.zeros(shape=(4, n))
matrix[0, :] = trace[n:-1] # casual
matrix[1, :] = np.flipud(trace[1:n + 1]) # acausal
matrix[2, :] = matrix[0, :] + matrix[1, :] # symmetric
matrix[3, :] = time[n:-1] # time
return matrix
def compute_max(trace, start, stop):
max_index = np.argmax(trace)
if max_index == (start - 1) or max_index == (stop - 1):
local_index = argrelextrema(trace, np.greater)[0]
local_value = trace[local_index]
if len(local_index) > 1:
sorted = np.argsort(local_value)[::-1]
max_index = sorted[1]
return max_index
# define temporal window of interest to isolate surface waves
start = int(distance / vmax * fs)
stop = int(distance / vmin * fs)
window_index = np.arange(start, stop)
window_length = stop - start
# matrices to store information on the group velocity and instantaneous period for plotting/saving
amplitude = np.zeros(shape=(3, window_length, len(periods)))
phase = np.zeros(shape=(3, window_length, len(periods)))
trace_label = ["P", "N", "S"]
# matrices to store group velocities and periods for plotting
group_velocities = np.zeros(shape=(3, len(periods)))
group_periods = np.zeros(shape=(3, len(periods)))
# to be filled with acceptable group velocities
output_matrix = np.zeros(shape=(3, len(periods))) * np.nan
# compute fft for the entire xcor just once
xfft = fft(xcor)
freqs = fftfreq(len(xcor), d=fs)
# freqs = np.linspace(0,fs,len(time))
# loop over periods
for k, T in enumerate(
periods): # enumerate to get index and value of interst
# construct Gaussian filter
filter_gaussian = np.exp(-((freqs - 1. / T) /
(alpha * 1. / T))**2) # Gaussian window
# multiply fft of signal by filter in the frequency domain (colvolution)
xfilter = xfft * filter_gaussian
# inverse Fourier transform back into time domain & keep only real components
xfreal = ifft(xfilter).real
# now check SNR for given period using Gaussian windowed trace
snr_period = compute_snr_period(xfreal, distance, fs)
# skip to next period if snr too low for any component
if min(snr_period) < snr_threshold:
continue
# get instantaneous amplitudes (envelope) and frequencies
x_env, x_phase = compute_instantaneous(xfreal, fs)
env_3 = split_trace(x_env,
time) # matrix of causal, acausal, symmetric parts
per_3 = split_trace(x_phase, time) # matrix of phase info for 3 traces
# make sure stations aren't too far apart or too close
if stop > len(env_3[0]) or start < 1:
return False
for j in range(3): #loop over causal, acausal, symmetric components
# extract tempoeral window of interest for plotting
# this copies the entries in env_3[j] with indices in [window_index]
# to the middle index in amplitude
amplitude[j, :, k] = env_3[j][window_index]
# extract tempoeral window of interest for disperion curve
phase[j, :, k] = per_3[j][window_index]
# get group velocity information
env_3[j][0:start] = 0.
env_3[j][stop:] = 0.
max_index = compute_max(env_3[j], start, stop)
travel_time = env_3[3][max_index]
if travel_time == 0:
gvel = np.nan # group velocity measurement
gper = np.nan # corresponding phase
else:
total_time = travel_time
gvel = distance / total_time
# save group velocity and group period
group_velocities[j][k] = gvel
group_periods[j][k] = T
# prune based on wavelength
for i in range(3):
for period_index, T in enumerate(periods):
velocity = group_velocities[i, period_index]
if velocity == 0:
continue
# check min number of wavelengths
wavelength = velocity * T
if not wavelength > 0:
continue
n_waves = int(distance / wavelength)
if n_waves < 2:
continue
output_matrix[i, period_index] = round(velocity, 4)
# prune based on mismatch between causal and acausal measurements
for i in range(len(periods)):
difference = abs(output_matrix[0, i] - output_matrix[1, i])
if difference < 0.2:
continue
else:
output_matrix[:, i] = np.nan
# prune based on maximum velocity
for i in range(len(periods)):
if max(output_matrix[:, i]) >= vmax:
output_matrix[:, i] = np.nan
# return valid group velocity measurements
return output_matrix
import matplotlib.pyplot as plt
import numpy as np
from numpy.fft import fft, fftfreq, ifft
from numpy import cosh as cosh
from matplotlib.animation import FuncAnimation
N = 400
length = 100
dx = length / N
x = dx * np.arange(N)
k = 2 * np.pi / length * fftfreq(N) * N
#initial soliton solution profile
beta = 0.5
soliton_i = 0.5 * beta / (cosh((beta)**0.5 / 2 * (x - length / 2)))**2
#soliton_i = beta/2 * ( cosh( sqrt(beta)/2 * (x-length/2) ) )**(-2.)
L_ = -(1j * k)**3
def N_(u):
return -6 * 0.5 * 1j * k * fft(u**2)
#time step
dt = 0.05
v2 = plt.subplot(2,1,2)
v2.scatter(x1,y1)
y = interp.interp1d(x1,y1, kind = "cubic")
x = np.linspace(0,6,1000)
v2.plot(x,y(x), c="r")
plt.legend(["datos", "cubica"])
plt.savefig("interp.png")
n=100
x= np.linspace(0,2*3.14,n)
y = np.sin(x) +np.random.rand(n)
yo = f.fftshift( f.fft(y) )
xo = f.fftshift( f.fftfreq(n) )
plt.figure()
v1 =plt.subplot(4,1,1)
v1.plot(xo,abs(yo),label="fourier")
plt.legend()
v2 =plt.subplot(4,1,2)
v2.plot(x,y,label="funcion")
plt.legend()
yo[abs(xo) > 0.01] = 0
#delta = np.where(abs(xo) > 0)
v3 =plt.subplot(4,1,3)
plt.xlim(-1.0,1.0)
'../ML/data/freq.wav') # ./ML 当前目录的ML文件夹， ../ML 上一级目录的ML文件夹
print(sample_rate) # 44100 采样率
print(sigs.shape) # (132300,) 单声道 一列
print('time_wav = ', sigs.shape[0] / sample_rate, 's') # 音频播放时长
print(sigs[:5]) # [ -3893 -18346 -10040 19031 30756] / 2**15 == 声场强度
sigs = sigs / 2**15
print(sigs)
#[-0.11880493 -0.55987549 -0.30639648 0.58078003 0.93859863 0.17752075
# -0.40713501 -0.37188721 0.23010254 0.79000854 0.60708618 -0.33999634
# -0.45401001 -0.03741455 0.72381592 0.71847534 -0.00985718 -0.65844727
# -0.23483276 0.56091309 0.84158325 0.03463745 -0.5508728 -0.49267578
# 0.39089966 0.89395142 0.54898071 -0.48608398 -0.44085693 0.05752563]
times = np.arange(sigs.shape[0]) / sample_rate
freqs = nf.fftfreq(
len(sigs),
d=1 / sample_rate # 采样周期 = 1/ 采样频率
)
ffts = nf.fft(sigs) # 快速傅里叶变换
pows = np.abs(ffts)
plt.figure('Audio_Signal', facecolor='gray')
plt.title('Audio_Signal')
plt.xlabel('Time')
plt.ylabel('Signal')
plt.tick_params(labelsize=10)
plt.grid(linestyle=':')
plt.plot(times, sigs, '-', c='purple', label='Signal')
plt.legend()
plt.figure('Audio_Frequency', facecolor='gray')
plt.title('Audio_Frequency')
IMU2_X_gt[i] = 0
if IMU2_Y[i]>-thresh and IMU2_Y[i]<thresh:
IMU2_Y_gt[i] = 0
if IMU2_Z[i]>-thresh and IMU2_Z[i]<thresh:
IMU2_Z_gt[i] = 0
IMU1_X_gt = np.array(IMU1_X_gt)
IMU1_Y_gt = np.array(IMU1_Y_gt)
IMU1_Z_gt = np.array(IMU1_Z_gt)
IMU2_X_gt = np.array(IMU2_X_gt)
IMU2_Y_gt = np.array(IMU2_Y_gt)
IMU2_Z_gt = np.array(IMU2_Z_gt)
#-------FFT -------------
n =len(IMU1_X_raw)
freqs = fftfreq(n)
idx = np.argsort(freqs)
mask = freqs > 0
#fft values
fft_IMU1_X_raw = fft(IMU1_X_raw)
fft_IMU1_Y_raw = fft(IMU1_Y_raw)
fft_IMU1_Z_raw = fft(IMU1_Z_raw)
fft_IMU2_X_raw = fft(IMU2_X_raw)
fft_IMU2_Y_raw = fft(IMU2_Y_raw)
fft_IMU2_Z_raw = fft(IMU2_Z_raw)
# true theorical fft
fft_th_IMU1_X_raw = 2.0*np.abs(fft_IMU1_X_raw)
fft_th_IMU1_Y_raw = 2.0*np.abs(fft_IMU1_Y_raw)
fft_th_IMU1_Z_raw = 2.0*np.abs(fft_IMU1_Z_raw)
fft_th_IMU2_X_raw = 2.0*np.abs(fft_IMU2_X_raw)
fft_th_IMU2_Y_raw = 2.0*np.abs(fft_IMU2_Y_raw)
def fourier_fit(x,
y,
n_predict=0,
x_smooth=None,
n_pts=n_pts_smooth,
n_harm=default_fourier_n_harm):
"""
Creates a Fourier fit of a NumPy array. Also supports extrapolation.
Credit goes to https://gist.github.com/tartakynov/83f3cd8f44208a1856ce.
Parameters
----------
x, y: numpy.ndarray
1D NumPy arrays of the x and y values to fit to.
Must not contain NaNs.
n_predict: int
The number of points to extrapolate.
The points will be spaced evenly by the mean spacing of values in `x`.
x_smooth: list-like, optional
The exact x values to interpolate for. Supercedes `n_pts`.
n_pts: int, optional
The number of evenly spaced points spanning the range of `x` to interpolate for.
n_harm: int
The number of harmonics to use. A higher value yields a closer fit.
Returns
-------
x_smooth, y_smooth: numpy.ndarray
The smoothed x and y values of the curve fit.
"""
if x_smooth is None:
x_smooth_inds = np.linspace(0, len(x) - 1, n_pts)
x_smooth = np.interp(x_smooth_inds, np.arange(len(x)), x)
n_predict_smooth = int((len(x_smooth) / len(x)) * n_predict)
# These points are evenly spaced for the fourier fit implementation we use.
# More points are selected than are in `x_smooth` so we can interpolate accurately.
fourier_mult_pts = 2
x_smooth_fourier = np.linspace(x_smooth.min(), x_smooth.max(),
fourier_mult_pts * len(x_smooth))
y_smooth_fourier = np.interp(x_smooth_fourier, x, y)
n_predict_smooth_fourier = int(
(len(x_smooth_fourier) / len(x)) * n_predict)
# Perform the Fourier fit and extrapolation.
n = y_smooth_fourier.size
t = np.arange(0, n)
p = np.polyfit(t, y_smooth_fourier, 1) # find linear trend in arr
x_notrend = y_smooth_fourier - p[0] * t # detrended arr
x_freqdom = fft.fft(x_notrend) # detrended arr in frequency domain
f = fft.fftfreq(n) # frequencies
# sort indexes by frequency, lower -> higher
indexes = list(range(n))
indexes.sort(key=lambda i: np.absolute(x_freqdom[i]))
indexes.reverse()
t = np.arange(0, n + n_predict_smooth_fourier)
restored_sig = np.zeros(t.size)
for i in indexes[:1 + n_harm * 2]:
ampli = np.absolute(x_freqdom[i]) / n # amplitude
phase = np.angle(x_freqdom[i]) # phase
restored_sig += ampli * np.cos(2 * np.pi * f[i] * t + phase)
y_smooth_fourier = restored_sig + p[0] * t
# Find the points in `x_smooth_fourier` that are near to points in `x_smooth`
# and then interpolate the y values to match the new x values.
x_smooth = x_smooth_fourier[np.searchsorted(x_smooth_fourier, x_smooth)]
# Ensure `x_smooth` includes the extrapolations.
mean_x_smooth_space = np.diff(x_smooth).mean()
x_predict_smooth = np.linspace(
x_smooth[-1] + mean_x_smooth_space,
x_smooth[-1] + mean_x_smooth_space * n_predict_smooth,
n_predict_smooth)
x_smooth = np.concatenate((x_smooth, x_predict_smooth))
# Ensure `x_smooth_fourier` includes the extrapolations.
mean_x_smooth_fourier_space = np.diff(x_smooth).mean()
x_predict_smooth_fourier = \
np.linspace(
x_smooth_fourier[-1] + mean_x_smooth_fourier_space,
x_smooth_fourier[-1] + mean_x_smooth_fourier_space * n_predict_smooth_fourier,
n_predict_smooth_fourier)
x_smooth_fourier = np.concatenate(
(x_smooth_fourier, x_predict_smooth_fourier))
y_smooth = np.interp(x_smooth, x_smooth_fourier, y_smooth_fourier)
return x_smooth, y_smooth
def __init__(self, **kwargs):
"""
The following parameters must be specified
X_gridDIM - the coordinate grid size
X_amplitude - maximum value of the coordinates
P_gridDIM - the momentum grid size
P_amplitude - maximum value of the momentum
V(x) - potential energy (as a function)
K(p) - momentum dependent part of the hamiltonian (as a function)
kT - the temperature (if kT = 0, then the ground state Wigner function will be obtained.)
dbeta - (optional) 1/kT step size
"""
# save all attributes
for name, value in kwargs.items():
setattr(self, name, value)
# Check that all attributes were specified
try:
self.X_gridDIM
except AttributeError:
raise AttributeError(
"Coordinate grid size (X_gridDIM) was not specified")
assert self.X_gridDIM % 2 == 0, "Coordinate grid size (X_gridDIM) must be even"
try:
self.P_gridDIM
except AttributeError:
raise AttributeError(
"Momentum grid size (P_gridDIM) was not specified")
assert self.P_gridDIM % 2 == 0, "Momentum grid size (P_gridDIM) must be even"
try:
self.X_amplitude
except AttributeError:
raise AttributeError(
"Coordinate grid range (X_amplitude) was not specified")
try:
self.P_amplitude
except AttributeError:
raise AttributeError(
"Momentum grid range (P_amplitude) was not specified")
try:
self.V
except AttributeError:
raise AttributeError("Potential energy (V) was not specified")
try:
self.K
except AttributeError:
raise AttributeError("Momentum dependence (K) was not specified")
try:
self.kT
except AttributeError:
raise AttributeError("Temperature (kT) was not specified")
if self.kT > 0:
try:
self.dbeta
except AttributeError:
# if dbeta is not defined, just choose some value
self.dbeta = 0.01
# get number of dbeta steps to reach the desired Gibbs state
self.num_beta_steps = 1. / (self.kT * self.dbeta)
if round(self.num_beta_steps) != self.num_beta_steps:
# Changing self.dbeta so that num_beta_steps is an exact integer
self.num_beta_steps = round(self.num_beta_steps)
self.dbeta = 1. / (self.kT * self.num_beta_steps)
self.num_beta_steps = int(self.num_beta_steps)
else:
raise NotImplemented(
"The calculation of the ground state Wigner function has not been implemnted"
)
###################################################################################
#
# Generate grids
#
###################################################################################
# get coordinate and momentum step sizes
self.dX = 2. * self.X_amplitude / self.X_gridDIM
self.dP = 2. * self.P_amplitude / self.P_gridDIM
# coordinate grid
self.X = np.linspace(-self.X_amplitude, self.X_amplitude - self.dX,
self.X_gridDIM)
self.X = self.X[np.newaxis, :]
# Lambda grid (variable conjugate to the coordinate)
self.Lambda = fft.fftfreq(self.X_gridDIM, self.dX / (2 * np.pi))
# take only first half, as required by the real fft
self.Lambda = self.Lambda[:(1 + self.X_gridDIM // 2)]
#
self.Lambda = self.Lambda[np.newaxis, :]
# momentum grid
self.P = np.linspace(-self.P_amplitude, self.P_amplitude - self.dP,
self.P_gridDIM)
self.P = self.P[:, np.newaxis]
# Theta grid (variable conjugate to the momentum)
self.Theta = fft.fftfreq(self.P_gridDIM, self.dP / (2 * np.pi))
# take only first half, as required by the real fft
self.Theta = self.Theta[:(1 + self.P_gridDIM // 2)]
#
self.Theta = self.Theta[:, np.newaxis]
###################################################################################
#
# Pre-calculate exponents used for the split operator propagation
#
###################################################################################
# Get the sum of the potential energy contributions
self.expV = self.V(self.X - 0.5 * self.Theta) + self.V(self.X + 0.5 *
self.Theta)
self.expV *= -0.25 * self.dbeta
# Make sure that the largest value is zero
self.expV -= self.expV.max()
# such that the following exponent is never larger then one
np.exp(self.expV, out=self.expV)
# Get the sum of the kinetic energy contributions
self.expK = self.K(self.P + 0.5 * self.Lambda) + self.K(self.P - 0.5 *
self.Lambda)
self.expK *= -0.5 * self.dbeta
# Make sure that the largest value is zero
self.expK -= self.expK.max()
# such that the following exponent is never larger then one
np.exp(self.expK, out=self.expK)
