Coverage for gpaw/test/maths/test_fftw.py: 100%

28 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-07-14 00:18 +0000

1import time 

2 

3import pytest 

4 

5import gpaw.fftw as fftw 

6 

7N = 1 

8# n1, n2, n3 = 32, 28, 128 

9n1, n2, n3 = 8, 6, 12 

10 

11 

12def check(Plan, flags, input, output, sign): 

13 t0 = time.time() 

14 plan = Plan(input, output, sign, flags) 

15 t1 = time.time() 

16 t = 0.0 

17 for i in range(N): 

18 input[:] = 1.3 

19 t2 = time.time() 

20 plan.execute() 

21 t3 = time.time() 

22 t += t3 - t2 

23 return t1 - t0, t / N 

24 

25 

26@pytest.mark.skipif(not fftw.have_fftw(), reason='No FFTW') 

27def test_fft(): 

28 a1 = fftw.empty((n1, n2, n3), complex) 

29 a2 = fftw.empty((n1, n2, n3), complex) 

30 b = fftw.empty((n1, n2, n3 // 2 + 1), complex) 

31 c1 = b.view(dtype=float)[:, :, :n3] 

32 c2 = fftw.empty((n1, n2, n3 // 2), complex).view(dtype=float) 

33 for input, output, sign in [ 

34 (a1, a1, -1), 

35 (a1, a2, -1), 

36 (b, c1, 1), 

37 (b, c2, 1), 

38 (c1, b, -1), 

39 (c2, b, -1)]: 

40 for Plan, flags in [(fftw.NumpyFFTPlan, 117), 

41 (fftw.FFTWPlan, fftw.ESTIMATE), 

42 (fftw.FFTWPlan, fftw.MEASURE), 

43 (fftw.FFTWPlan, fftw.PATIENT), 

44 (fftw.FFTWPlan, fftw.EXHAUSTIVE)]: 

45 tplan, tfft = check(Plan, flags, input, output, sign) 

46 print(('%-12s %3d %10.6f %10.6f' % 

47 (Plan.__name__, flags, tplan, tfft)))