Coverage for gpaw/test/new/test_move_wfs.py: 100%

1import numpy as np 

2import pytest 

3from gpaw.core import PWDesc 

4from gpaw.gpu import cupy as cp 

5from gpaw.mpi import world 

6from gpaw.new.pwfd.move_wfs import move_wave_functions 

7from gpaw.setup import create_setup 

8from gpaw.new.c import GPU_AWARE_MPI 

9from gpaw.gpu.mpi import CuPyMPI 

10from gpaw.core.atom_arrays import AtomArraysLayout 

11 

12 

13@pytest.mark.parametrize('xp', 

14 [np, 

15 pytest.param(cp, marks=pytest.mark.gpu)]) 

16def test_move(xp): 

17 comm = world if GPU_AWARE_MPI else CuPyMPI(world) 

18 pw = PWDesc(ecut=25, cell=[2, 2, 2], kpt=[0.25, 0.25, 0.0], comm=comm) 

19 psit_nG = pw.zeros(2, xp=xp) 

20 print(psit_nG.data.shape) 

21 pos1 = np.zeros((1, 3)) 

22 pos2 = pos1 + 0.1 

23 setup = create_setup('H') 

24 P_ani = AtomArraysLayout([(5,)], comm, complex, xp).empty(2) 

25 if comm.rank == 0: 

26 P_ani[0][:] = 1 + 2j 

27 move_wave_functions(pos1, pos2, P_ani, psit_nG, [setup]) 

28 move_wave_functions(pos2, pos1, P_ani, psit_nG, [setup]) 

29 assert abs(psit_nG.to_xp(np).data).max() < 1e-12 

30 

31 

32test_move(np)