def test_pdos_1d():
pad=lambda x: pad_zeros(x, nadd=len(x)-1)
n=500; w=welch(n)
# 1 second signal
t=np.linspace(0,1,n); dt=t[1]-t[0]
# sum of sin()s with random freq and phase shift, 10 frequencies from
# f=0...100 Hz
v=np.array([np.sin(2*np.pi*f*t + rand()*2*np.pi) for f in rand(10)*100]).sum(0)
f=np.fft.fftfreq(2*n-1, dt)[:n]
c1=mirror(ifft(abs(fft(pad(v)))**2.0)[:n].real)
c2=correlate(v,v,'full')
c3=mirror(acorr(v,norm=False))
assert np.allclose(c1, c2)
assert np.allclose(c1, c3)
p1=(abs(fft(pad(v)))**2.0)[:n]
p2=(abs(fft(mirror(acorr(v,norm=False)))))[:n]
assert np.allclose(p1, p2)
p1=(abs(fft(pad(v*w)))**2.0)[:n]
p2=(abs(fft(mirror(acorr(v*w,norm=False)))))[:n]
assert np.allclose(p1, p2)
def test_pdos_1d():
pad=lambda x: pad_zeros(x, nadd=len(x)-1)
n=500; w=welch(n)
# 1 second signal
t=np.linspace(0,1,n); dt=t[1]-t[0]
# sum of sin()s with random freq and phase shift, 10 frequencies from
# f=0...100 Hz
v=np.array([np.sin(2*np.pi*f*t + rand()*2*np.pi) for f in rand(10)*100]).sum(0)
f=np.fft.fftfreq(2*n-1, dt)[:n]
c1=mirror(ifft(abs(fft(pad(v)))**2.0)[:n].real)
c2=correlate(v,v,'full')
c3=mirror(acorr(v,norm=False))
assert np.allclose(c1, c2)
assert np.allclose(c1, c3)
p1=(abs(fft(pad(v)))**2.0)[:n]
p2=(abs(fft(mirror(acorr(v,norm=False)))))[:n]
assert np.allclose(p1, p2)
p1=(abs(fft(pad(v*w)))**2.0)[:n]
p2=(abs(fft(mirror(acorr(v*w,norm=False)))))[:n]
assert np.allclose(p1, p2)
plots = mpl.prepare_plots(["freq", "filt_pad", "filt_nopad"])
nyq = 100 # Hz
df = 1.0 # Hz
dt, nstep = fftsample(nyq, df, mode="f")
t = np.linspace(0, 1, int(nstep))
filt1 = FIRFilter(cutoff=[10, 50], nyq=nyq, mode="bandpass", ripple=60, width=10)
filt2 = FIRFilter(cutoff=[10, 50], nyq=nyq, mode="bandpass", ntaps=100, window="hamming")
plots["freq"].ax.plot(filt1.w, abs(filt1.h), label="filt1")
plots["freq"].ax.plot(filt2.w, abs(filt2.h), label="filt2")
plots["freq"].ax.legend()
for pad in [True, False]:
x = np.sin(2 * pi * 20 * t) + np.sin(2 * pi * 80 * t)
if pad:
x = pad_zeros(x, nadd=len(x))
pl = plots["filt_pad"]
else:
pl = plots["filt_nopad"]
f = np.fft.fftfreq(len(x), dt)
sl = slice(0, len(x) / 2, None)
win = hanning(len(x))
pl.ax.plot(f[sl], np.abs(fft(x)[sl]), label="fft(x)")
pl.ax.plot(f[sl], np.abs(fft(filt1(x))[sl]), label="fft(filt1(x))")
pl.ax.plot(f[sl], np.abs(fft(filt1(win * x))[sl]), label="fft(filt1(hanning*x))")
pl.ax.plot(f[sl], np.abs(fft(filt2(win * x))[sl]), label="fft(filt2(hanning*x))")
pl.ax.set_title("zero pad = %s" % pad)
pl.ax.legend()
plt.show()
mode='bandpass',
ripple=60,
width=10)
filt2 = FIRFilter(cutoff=[10, 50],
nyq=nyq,
mode='bandpass',
ntaps=100,
window='hamming')
plots['freq'].ax.plot(filt1.w, abs(filt1.h), label='filt1')
plots['freq'].ax.plot(filt2.w, abs(filt2.h), label='filt2')
plots['freq'].ax.legend()
for pad in [True, False]:
x = np.sin(2 * pi * 20 * t) + np.sin(2 * pi * 80 * t)
if pad:
x = pad_zeros(x, nadd=len(x))
pl = plots['filt_pad']
else:
pl = plots['filt_nopad']
f = np.fft.fftfreq(len(x), dt)
sl = slice(0, len(x) / 2, None)
win = hanning(len(x))
pl.ax.plot(f[sl], np.abs(fft(x)[sl]), label='fft(x)')
pl.ax.plot(f[sl], np.abs(fft(filt1(x))[sl]), label='fft(filt1(x))')
pl.ax.plot(f[sl],
np.abs(fft(filt1(win * x))[sl]),
label='fft(filt1(hanning*x))')
pl.ax.plot(f[sl],
np.abs(fft(filt2(win * x))[sl]),
label='fft(filt2(hanning*x))')
pl.ax.set_title('zero pad = %s' % pad)
def pdos(vel,
dt=1.0,
m=None,
full_out=False,
area=1.0,
window=True,
npad=None,
tonext=False,
mirr=False,
method='direct'):
"""Phonon DOS by FFT of the VACF or direct FFT of atomic velocities.
Integral area is normalized to `area`. It is possible (and recommended) to
zero-padd the velocities (see `npad`).
Parameters
----------
vel : 3d array (nstep, natoms, 3)
atomic velocities
dt : time step
m : 1d array (natoms,),
atomic mass array, if None then mass=1.0 for all atoms is used
full_out : bool
area : float
normalize area under frequency-PDOS curve to this value
window : bool
use Welch windowing on data before FFT (reduces leaking effect,
recommended)
npad : {None, int}
method='direct' only: Length of zero padding along `axis`. `npad=None`
= no padding, `npad > 0` = pad by a length of ``(nstep-1)*npad``. `npad
> 5` usually results in sufficient interpolation.
tonext : bool
method='direct' only: Pad `vel` with zeros along `axis` up to the next
power of two after the array length determined by `npad`. This gives
you speed, but variable (better) frequency resolution.
mirr : bool
method='vacf' only: mirror one-sided VACF at t=0 before fft
Returns
-------
if full_out = False
| ``(faxis, pdos)``
| faxis : 1d array [1/unit(dt)]
| pdos : 1d array, the phonon DOS, normalized to `area`
if full_out = True
| if method == 'direct':
| ``(faxis, pdos, (full_faxis, full_pdos, split_idx))``
| if method == 'vavcf':
| ``(faxis, pdos, (full_faxis, full_pdos, split_idx, vacf, fft_vacf))``
| fft_vacf : 1d complex array, result of fft(vacf) or fft(mirror(vacf))
| vacf : 1d array, the VACF
Examples
--------
>>> from pwtools.constants import fs,rcm_to_Hz
>>> tr = Trajectory(...)
>>> # freq in [Hz] if timestep in [s]
>>> freq,dos = pdos(tr.velocity, m=tr.mass, dt=tr.timestep*fs,
>>> method='direct', npad=1)
>>> # frequency in [1/cm]
>>> plot(freq/rcm_to_Hz, dos)
Notes
-----
padding (only method='direct'): With `npad` we pad the velocities `vel`
with ``npad*(nstep-1)`` zeros along `axis` (the time axis) before FFT
b/c the signal is not periodic. For `npad=1`, this gives us the exact
same spectrum and frequency resolution as with ``pdos(...,
method='vacf',mirr=True)`` b/c the array to be fft'ed has length
``2*nstep-1`` along the time axis in both cases (remember that the
array length = length of the time axis influences the freq.
resolution). FFT is only fast for arrays with length = a power of two.
Therefore, you may get very different fft speeds depending on whether
``2*nstep-1`` is a power of two or not (in most cases it won't). Try
using `tonext` but remember that you get another (better) frequency
resolution.
References
----------
[1] Phys Rev B 47(9) 4863, 1993
See Also
--------
:func:`pwtools.signal.fftsample`
:func:`pwtools.signal.acorr`
:func:`direct_pdos`
:func:`vacf_pdos`
"""
mass = m
# assume vel.shape = (nstep,natoms,3)
axis = 0
assert vel.shape[-1] == 3
if mass is not None:
assert len(mass) == vel.shape[1], "len(mass) != vel.shape[1]"
# define here b/c may be used twice below
mass_bc = mass[None, :, None]
if window:
sl = [None] * vel.ndim
sl[axis] = slice(None) # ':'
vel2 = vel * (welch(vel.shape[axis])[sl])
else:
vel2 = vel
# handle options which are mutually exclusive
if method == 'vacf':
assert npad in [0, None], "use npad={0,None} for method='vacf'"
# padding
if npad is not None:
nadd = (vel2.shape[axis] - 1) * npad
if tonext:
vel2 = pad_zeros(vel2,
tonext=True,
tonext_min=vel2.shape[axis] + nadd,
axis=axis)
else:
vel2 = pad_zeros(vel2, tonext=False, nadd=nadd, axis=axis)
if method == 'direct':
full_fft_vel = np.abs(fft(vel2, axis=axis))**2.0
full_faxis = np.fft.fftfreq(vel2.shape[axis], dt)
split_idx = len(full_faxis) // 2
faxis = full_faxis[:split_idx]
# First split the array, then multiply by `mass` and average. If
# full_out, then we need full_fft_vel below, so copy before slicing.
arr = full_fft_vel.copy() if full_out else full_fft_vel
fft_vel = num.slicetake(arr,
slice(0, split_idx),
axis=axis,
copy=False)
if mass is not None:
fft_vel *= mass_bc
# average remaining axes, summing is enough b/c normalization is done below
# sums: (nstep, natoms, 3) -> (nstep, natoms) -> (nstep,)
pdos = num.sum(fft_vel, axis=axis, keepdims=True)
default_out = (faxis, num.norm_int(pdos, faxis, area=area))
if full_out:
# have to re-calculate this here b/c we never calculate the full_pdos
# normally
if mass is not None:
full_fft_vel *= mass_bc
full_pdos = num.sum(full_fft_vel, axis=axis, keepdims=True)
extra_out = (full_faxis, full_pdos, split_idx)
return default_out + extra_out
else:
return default_out
elif method == 'vacf':
vacf = fvacf(vel2, m=mass)
if mirr:
fft_vacf = fft(mirror(vacf))
else:
fft_vacf = fft(vacf)
full_faxis = np.fft.fftfreq(fft_vacf.shape[axis], dt)
full_pdos = np.abs(fft_vacf)
split_idx = len(full_faxis) // 2
faxis = full_faxis[:split_idx]
pdos = full_pdos[:split_idx]
default_out = (faxis, num.norm_int(pdos, faxis, area=area))
extra_out = (full_faxis, full_pdos, split_idx, vacf, fft_vacf)
if full_out:
return default_out + extra_out
else:
return default_out
coords = np.sin(2 * pi * freqs[:, None] * taxis).sum(axis=0)
# below, `arr` is used fo all following calcs
##arr = coords
arr = np.diff(coords) # "velocity"
# Our methods 1 and 2 must result in exactly the same after
# norm_int()'ing them (within numerical noise).
# In both, we use a Welch window as in fourier.x .
# 1) Zero-pad arr at the end and fft directly. For padding, use nadd=N-1 to get
# exactly the same signal length as y2 b/c in case of y2: mirror(rand(N)).shape
# = (2N-1,), i.e. the point at t=0 apprears only once. Padding increases the
# frequency resolution (almost factor 2) to exactly the same df as we get for
# method 2 b/c of the mirroring of the signal there.
fft_arr = signal.pad_zeros(arr * signal.welch(arr.shape[0]),
nadd=arr.shape[0] - 1)
y1 = np.abs(fft(fft_arr))**2
# 2) fft the autocorrelation of `arr`
fft_arr = pydos.mirror(signal.acorr(arr * signal.welch(arr.shape[0]),
method=5))
y2 = np.abs(fft(fft_arr))
# 3) fourier.x
#
# For the 1d case, we write the time trace in a format suggested in the CPMD
# manual and the fourier.x README file: awk '{ print $1, 0.0, 0.0, 0.0, 0.0,
# 0.0, $2; }' ENERGIES > ekinc.dat where ENERGIES is a CPMD output file. From
# that, only column 1 (time step) and some energy value from column 2 is used.
if use_fourier:
fourier_in_data = np.zeros((arr.shape[0], 7))
def pdos(vel, dt=1.0, m=None, full_out=False, area=1.0, window=True,
npad=None, tonext=False, mirr=False, method='direct'):
"""Phonon DOS by FFT of the VACF or direct FFT of atomic velocities.
Integral area is normalized to `area`. It is possible (and recommended) to
zero-padd the velocities (see `npad`).
Parameters
----------
vel : 3d array (nstep, natoms, 3)
atomic velocities
dt : time step
m : 1d array (natoms,),
atomic mass array, if None then mass=1.0 for all atoms is used
full_out : bool
area : float
normalize area under frequency-PDOS curve to this value
window : bool
use Welch windowing on data before FFT (reduces leaking effect,
recommended)
npad : {None, int}
method='direct' only: Length of zero padding along `axis`. `npad=None`
= no padding, `npad > 0` = pad by a length of ``(nstep-1)*npad``. `npad
> 5` usually results in sufficient interpolation.
tonext : bool
method='direct' only: Pad `vel` with zeros along `axis` up to the next
power of two after the array length determined by `npad`. This gives
you speed, but variable (better) frequency resolution.
mirr : bool
method='vacf' only: mirror one-sided VACF at t=0 before fft
Returns
-------
if full_out = False
| ``(faxis, pdos)``
| faxis : 1d array [1/unit(dt)]
| pdos : 1d array, the phonon DOS, normalized to `area`
if full_out = True
| if method == 'direct':
| ``(faxis, pdos, (full_faxis, full_pdos, split_idx))``
| if method == 'vavcf':
| ``(faxis, pdos, (full_faxis, full_pdos, split_idx, vacf, fft_vacf))``
| fft_vacf : 1d complex array, result of fft(vacf) or fft(mirror(vacf))
| vacf : 1d array, the VACF
Examples
--------
>>> from pwtools.constants import fs,rcm_to_Hz
>>> tr = Trajectory(...)
>>> # freq in [Hz] if timestep in [s]
>>> freq,dos = pdos(tr.velocity, m=tr.mass, dt=tr.timestep*fs,
>>> method='direct', npad=1)
>>> # frequency in [1/cm]
>>> plot(freq/rcm_to_Hz, dos)
See Also
--------
:func:`direct_pdos`
:func:`vacf_pdos`
Notes
-----
padding (only method='direct'): With `npad` we pad the velocities `vel`
with ``npad*(nstep-1)`` zeros along `axis` (the time axis) before FFT
b/c the signal is not periodic. For `npad=1`, this gives us the exact
same spectrum and frequency resolution as with ``pdos(...,
method='vacf',mirr=True)`` b/c the array to be fft'ed has length
``2*nstep-1`` along the time axis in both cases (remember that the
array length = length of the time axis influences the freq.
resolution). FFT is only fast for arrays with length = a power of two.
Therefore, you may get very different fft speeds depending on whether
``2*nstep-1`` is a power of two or not (in most cases it won't). Try
using `tonext` but remember that you get another (better) frequency
resolution.
References
----------
[1] Phys Rev B 47(9) 4863, 1993
See Also
--------
pwtools.signal.fftsample
pwtools.signal.acorr
"""
mass = m
# assume vel.shape = (nstep,natoms,3)
axis = 0
assert vel.shape[-1] == 3
if mass is not None:
assert len(mass) == vel.shape[1], "len(mass) != vel.shape[1]"
# define here b/c may be used twice below
mass_bc = mass[None,:,None]
if window:
sl = [None]*vel.ndim
sl[axis] = slice(None) # ':'
vel2 = vel*(welch(vel.shape[axis])[sl])
else:
vel2 = vel
# handle options which are mutually exclusive
if method == 'vacf':
assert npad in [0,None], "use npad={0,None} for method='vacf'"
# padding
if npad is not None:
nadd = (vel2.shape[axis]-1)*npad
if tonext:
vel2 = pad_zeros(vel2, tonext=True,
tonext_min=vel2.shape[axis] + nadd,
axis=axis)
else:
vel2 = pad_zeros(vel2, tonext=False, nadd=nadd, axis=axis)
if method == 'direct':
full_fft_vel = np.abs(fft(vel2, axis=axis))**2.0
full_faxis = np.fft.fftfreq(vel2.shape[axis], dt)
split_idx = len(full_faxis)/2
faxis = full_faxis[:split_idx]
# First split the array, then multiply by `mass` and average. If
# full_out, then we need full_fft_vel below, so copy before slicing.
arr = full_fft_vel.copy() if full_out else full_fft_vel
fft_vel = num.slicetake(arr, slice(0, split_idx), axis=axis, copy=False)
if mass is not None:
fft_vel *= mass_bc
# average remaining axes, summing is enough b/c normalization is done below
# sums: (nstep, natoms, 3) -> (nstep, natoms) -> (nstep,)
pdos = num.sum(fft_vel, axis=axis, keepdims=True)
default_out = (faxis, num.norm_int(pdos, faxis, area=area))
if full_out:
# have to re-calculate this here b/c we never calculate the full_pdos
# normally
if mass is not None:
full_fft_vel *= mass_bc
full_pdos = num.sum(full_fft_vel, axis=axis, keepdims=True)
extra_out = (full_faxis, full_pdos, split_idx)
return default_out + extra_out
else:
return default_out
elif method == 'vacf':
vacf = fvacf(vel2, m=mass)
if mirr:
fft_vacf = fft(mirror(vacf))
else:
fft_vacf = fft(vacf)
full_faxis = np.fft.fftfreq(fft_vacf.shape[axis], dt)
full_pdos = np.abs(fft_vacf)
split_idx = len(full_faxis)/2
faxis = full_faxis[:split_idx]
pdos = full_pdos[:split_idx]
default_out = (faxis, num.norm_int(pdos, faxis, area=area))
extra_out = (full_faxis, full_pdos, split_idx, vacf, fft_vacf)
if full_out:
return default_out + extra_out
else:
return default_out
from math import pi
import numpy as np
from matplotlib import pyplot as plt
from pwtools import signal
from scipy.signal import convolve, gaussian, correlate
from scipy.fftpack import fft
nn = 200
nadd = 5*nn
t = np.linspace(0.123,0.567,nn)
x = np.sin(2*pi*10*t) + np.cos(2*pi*3*t) + np.sin(2*pi*30*t)
dt = t[1]-t[0]
pad_x = signal.pad_zeros(x, nadd=nadd)
pad_welch_x = signal.pad_zeros(x*signal.welch(nn), nadd=nadd)
kern = gaussian(M=20,std=2) # width M must be 6..10 x std
smooth_pad_x = convolve(signal.pad_zeros(x,nadd=nadd),kern,'same')/10.0
##mirr_x = signal.mirror(x)
##welch_mirr_x = signal.mirror(x)*signal.welch(2*nn-1)
##pad_welch_mirr_x = signal.pad_zeros(signal.mirror(x)*signal.welch(2*nn-1),
## nadd=2*nn-1)
plt.figure()
plt.plot(pad_x, label='pad_x (padded signal)')
plt.plot(pad_welch_x, label='pad_welch_x')
plt.plot(smooth_pad_x,label='smooth_pad_x')
plt.xlabel('time [s]')
plt.xlim(0,300)
plt.legend()
coords = np.sin(2*pi*freqs[:,None]*taxis).sum(axis=0)
# below, `arr` is used fo all following calcs
##arr = coords
arr = np.diff(coords) # "velocity"
# Our methods 1 and 2 must result in exactly the same after
# norm_int()'ing them (within numerical noise).
# In both, we use a Welch window as in fourier.x .
# 1) Zero-pad arr at the end and fft directly. For padding, use nadd=N-1 to get
# exactly the same signal length as y2 b/c in case of y2: mirror(rand(N)).shape
# = (2N-1,), i.e. the point at t=0 apprears only once. Padding increases the
# frequency resolution (almost factor 2) to exactly the same df as we get for
# method 2 b/c of the mirroring of the signal there.
fft_arr = signal.pad_zeros(arr*signal.welch(arr.shape[0]), nadd=arr.shape[0]-1)
y1 = np.abs(fft(fft_arr))**2
# 2) fft the autocorrelation of `arr`
fft_arr = pydos.mirror(signal.acorr(arr*signal.welch(arr.shape[0]), method=5))
y2 = np.abs(fft(fft_arr))
# 3) fourier.x
#
# For the 1d case, we write the time trace in a format suggested in the CPMD
# manual and the fourier.x README file: awk '{ print $1, 0.0, 0.0, 0.0, 0.0,
# 0.0, $2; }' ENERGIES > ekinc.dat where ENERGIES is a CPMD output file. From
# that, only column 1 (time step) and some energy value from column 2 is used.
if use_fourier:
fourier_in_data = np.zeros((arr.shape[0],7))
fourier_in_data[:,0] = np.arange(arr.shape[0])
# big, then the smoothing will filter out high frequencies.
from math import pi
import numpy as np
from matplotlib import pyplot as plt
from pwtools import signal
from scipy.signal import convolve, gaussian, correlate
from scipy.fftpack import fft
nn = 200
nadd = 5 * nn
t = np.linspace(0.123, 0.567, nn)
x = np.sin(2 * pi * 10 * t) + np.cos(2 * pi * 3 * t) + np.sin(2 * pi * 30 * t)
dt = t[1] - t[0]
pad_x = signal.pad_zeros(x, nadd=nadd)
pad_welch_x = signal.pad_zeros(x * signal.welch(nn), nadd=nadd)
kern = gaussian(M=20, std=2) # width M must be 6..10 x std
smooth_pad_x = convolve(signal.pad_zeros(x, nadd=nadd), kern, 'same') / 10.0
##mirr_x = signal.mirror(x)
##welch_mirr_x = signal.mirror(x)*signal.welch(2*nn-1)
##pad_welch_mirr_x = signal.pad_zeros(signal.mirror(x)*signal.welch(2*nn-1),
## nadd=2*nn-1)
plt.figure()
plt.plot(pad_x, label='pad_x (padded signal)')
plt.plot(pad_welch_x, label='pad_welch_x')
plt.plot(smooth_pad_x, label='smooth_pad_x')
plt.xlabel('time [s]')
plt.xlim(0, 300)
plt.legend()
