test_rfft.py (916B)
1 #!/usr/bin/env python3 2 # 3 # Copyright (c) 2021-2023 Xiaomi Corporation (authors: Fangjun Kuang) 4 5 6 import torch 7 8 import kaldi_native_fbank as knf 9 10 11 def test_rfft(): 12 N = 12 13 t = torch.arange(N) 14 r = torch.fft.rfft(t) 15 assert len(r) == N // 2 + 1, (len(r), N // 2 + 1) 16 17 real = r.real 18 imag = r.imag 19 print(r) 20 21 k = t.tolist() 22 rfft = knf.Rfft(N) 23 24 p = rfft.compute(k) 25 print(p) 26 return 27 28 assert abs(p[0] - real[0]) < 1e-5, (p[0], real[0]) 29 assert imag[0] == 0, imag[0] 30 31 assert abs(p[1] - real[-1]) < 1e-5, (p[1], real[-1]) 32 assert imag[-1] == 0, imag[-1] 33 34 for i in range(1, N // 2): 35 assert abs(p[2 * i] - real[i]) < 1e-5, (p[2 * i], real[i]) 36 # Note: the imaginary part is multiplied by negative 1 37 assert abs(p[2 * i + 1] - -1 * imag[i]) < 1e-5, (p[2 * i + 1], imag[i]) 38 39 40 def main(): 41 test_rfft() 42 43 44 if __name__ == "__main__": 45 main()