Coverage for gpaw/test/parallel/test_scalapack_diag_simple.py: 98%

1from time import time 

2 

3import numpy as np 

4from scipy.linalg import eigh 

5 

6# Set-up a simple matrix in parallel, diagonalize using ScaLAPACK 

7# D&C driver then compare *eigenvalues* with serial LAPACK diagonlize 

8 

9from gpaw.blacs import BlacsGrid 

10from gpaw.mpi import world 

11from gpaw.utilities.scalapack import scalapack_set, \ 

12 scalapack_zero 

13 

14 

15switch_uplo = {'U': 'L', 'L': 'U'} 

16rank = world.rank 

17tol = 5e-12 # eigenvalue tolerance 

18 

19 

20def test_parallel_scalapack_diag_simple(scalapack): 

21 nbands = 1000 

22 mb = 64 

23 

24 if world.size == 1: 

25 blacsshape = (1, 1) 

26 else: 

27 blacsshape = (world.size // 2, 2) 

28 

29 # Set-up BlacsGrud 

30 grid = BlacsGrid(world, *blacsshape) 

31 

32 # Create descriptor 

33 nndesc = grid.new_descriptor(nbands, nbands, mb, mb) 

34 H_nn = nndesc.empty(dtype=float) # outside BlacsGrid these are size zero 

35 C_nn = nndesc.empty(dtype=float) # outside BlacsGrid these are size zero 

36 eps_N = np.empty((nbands), dtype=float) # replicated on all MPI tasks 

37 # Fill ScaLAPACK array 

38 alpha = 0.1 # off-diagonal 

39 beta = 75.0 # diagonal 

40 uplo = 'L' # lower-triangular 

41 scalapack_set(nndesc, H_nn, alpha, beta, uplo) 

42 scalapack_zero(nndesc, H_nn, switch_uplo[uplo]) 

43 

44 t1 = time() 

45 # either interface will work, we recommend use the latter interface 

46 # scalapack_diagonalize_dc(nndesc, H_nn.copy(), C_nn, eps_N, 'L') 

47 nndesc.diagonalize_dc(H_nn.copy(), C_nn, eps_N) 

48 t2 = time() 

49 world.broadcast(eps_N, 0) # all MPI tasks now have eps_N 

50 world.barrier() # wait for everyone to finish 

51 

52 if rank == 0: 

53 print('ScaLAPACK diagonalize_dc', t2 - t1) 

54 

55 # Create replicated NumPy array 

56 diagonal = np.eye(nbands, dtype=float) 

57 offdiagonal = np.tril(np.ones((nbands, nbands)), -1) 

58 H0 = beta * diagonal + alpha * offdiagonal 

59 E0 = np.empty((nbands), dtype=float) 

60 

61 t1 = time() 

62 E0 = eigh(H0, eigvals_only=True) 

63 t2 = time() 

64 

65 if rank == 0: 

66 print('LAPACK diagonalize', t2 - t1) 

67 

68 delta = abs(E0 - eps_N).max() 

69 if rank == 0: 

70 print(delta) 

71 assert delta < tol