You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
108 lines
2.6 KiB
Python
108 lines
2.6 KiB
Python
import numpy as np
|
|
__all__ = ['dct', 'idct']
|
|
def dct(x, n=None):
|
|
"""
|
|
Discrete Cosine Transform
|
|
|
|
N-1
|
|
y[k] = 2* sum x[n]*cos(pi*k*(2n+1)/(2*N)), 0 <= k < N.
|
|
n=0
|
|
|
|
Examples
|
|
--------
|
|
>>> import numpy as np
|
|
>>> x = np.arange(5)
|
|
>>> np.abs(x-idct(dct(x)))<1e-14
|
|
array([ True, True, True, True, True], dtype=bool)
|
|
>>> np.abs(x-dct(idct(x)))<1e-14
|
|
array([ True, True, True, True, True], dtype=bool)
|
|
|
|
Reference
|
|
---------
|
|
http://en.wikipedia.org/wiki/Discrete_cosine_transform
|
|
http://users.ece.utexas.edu/~bevans/courses/ee381k/lectures/
|
|
"""
|
|
fft = np.fft.fft
|
|
x = np.atleast_1d(x)
|
|
|
|
if n is None:
|
|
n = x.shape[-1]
|
|
|
|
if x.shape[-1] < n:
|
|
n_shape = x.shape[:-1] + (n - x.shape[-1],)
|
|
xx = np.hstack((x, np.zeros(n_shape)))
|
|
else:
|
|
xx = x[..., :n]
|
|
|
|
real_x = np.all(np.isreal(xx))
|
|
if (real_x and (np.remainder(n, 2) == 0)):
|
|
xp = 2 * fft(np.hstack((xx[..., ::2], xx[..., ::-2])))
|
|
else:
|
|
xp = fft(np.hstack((xx, xx[..., ::-1])))
|
|
xp = xp[..., :n]
|
|
|
|
w = np.exp(-1j * np.arange(n) * np.pi / (2 * n))
|
|
|
|
y = xp * w
|
|
|
|
if real_x:
|
|
return y.real
|
|
else:
|
|
return y
|
|
|
|
def idct(x, n=None):
|
|
"""
|
|
Inverse Discrete Cosine Transform
|
|
|
|
N-1
|
|
x[k] = 1/N sum w[n]*y[n]*cos(pi*k*(2n+1)/(2*N)), 0 <= k < N.
|
|
n=0
|
|
|
|
w(0) = 1/2
|
|
w(n) = 1 for n>0
|
|
|
|
Examples
|
|
--------
|
|
>>> import numpy as np
|
|
>>> x = np.arange(5)
|
|
>>> np.abs(x-idct(dct(x)))<1e-14
|
|
array([ True, True, True, True, True], dtype=bool)
|
|
>>> np.abs(x-dct(idct(x)))<1e-14
|
|
array([ True, True, True, True, True], dtype=bool)
|
|
|
|
Reference
|
|
---------
|
|
http://en.wikipedia.org/wiki/Discrete_cosine_transform
|
|
http://users.ece.utexas.edu/~bevans/courses/ee381k/lectures/
|
|
"""
|
|
|
|
ifft = np.fft.ifft
|
|
x = np.atleast_1d(x)
|
|
|
|
if n is None:
|
|
n = x.shape[-1]
|
|
|
|
w = np.exp(1j * np.arange(n) * np.pi / (2 * n))
|
|
|
|
if x.shape[-1] < n:
|
|
n_shape = x.shape[:-1] + (n - x.shape[-1],)
|
|
xx = np.hstack((x, np.zeros(n_shape))) * w
|
|
else:
|
|
xx = x[..., :n] * w
|
|
|
|
real_x = np.all(np.isreal(x))
|
|
if (real_x and (np.remainder(n, 2) == 0)):
|
|
xx[..., 0] = xx[..., 0] * 0.5
|
|
yp = ifft(xx)
|
|
y = np.zeros(xx.shape, dtype=complex)
|
|
y[..., ::2] = yp[..., :n / 2]
|
|
y[..., ::-2] = yp[..., n / 2::]
|
|
else:
|
|
yp = ifft(np.hstack((xx, np.zeros_like(xx[..., 0]), np.conj(xx[..., :0:-1]))))
|
|
y = yp[..., :n]
|
|
|
|
if real_x:
|
|
return y.real
|
|
else:
|
|
return y
|