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()