Coverage for gpaw/test/gpu/test_precondition.py: 100%

1import numpy as np 

2from gpaw.core import PWDesc 

3from gpaw.gpu import cupy as cp, T 

4from gpaw.new.pw.hamiltonian import precondition 

5import pytest 

6 

7 

8@pytest.mark.gpu 

9@pytest.mark.parametrize('xp', [np, cp]) 

10def test_prec(xp): 

11 a = 2 

12 pw = PWDesc(cell=[a, a, a], ecut=200 / 27, dtype=complex) 

13 n = 2 

14 psit_nG, residual_nG, out_nG = pw.zeros((3, n), xp=xp) 

15 psit_nG.data[:, :2] = 1.0 

16 residual_nG.data[:] = 1.0 

17 print(residual_nG) 

18 for _ in range(1): 

19 with T(): 

20 precondition(psit_nG, residual_nG, out_nG)