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