Coverage for gpaw/test/test_diagonalizer_backend.py: 96%

1import numpy as np 

2import pytest 

3from ase.parallel import world 

4from gpaw.utilities import compiled_with_sl 

5from gpaw.utilities.elpa import LibElpa 

6from gpaw.eigensolvers.diagonalizerbackend import ( 

7 DistributedBlacsDiagonalizer, 

8 ScipyDiagonalizer, 

9 ScalapackDiagonalizer, 

10 ElpaDiagonalizer) 

11 

12 

13def prepare_eigensolver_matrices(size_of_matrices, dtype): 

14 

15 matrix_dimensions = [size_of_matrices, size_of_matrices] 

16 rng = np.random.Generator(np.random.PCG64(24589246)) 

17 A = rng.random(matrix_dimensions).astype(dtype) 

18 B = rng.random(matrix_dimensions).astype(dtype) 

19 

20 if dtype == complex: 

21 A += 1j * rng.random(matrix_dimensions) 

22 B += 1j * rng.random(matrix_dimensions) 

23 A = A + A.T.conj() 

24 B = B + B.T.conj() 

25 # Make sure B is positive definite 

26 B += np.eye(size_of_matrices) * size_of_matrices 

27 

28 return A, B 

29 

30 

31@pytest.fixture(params=['eigh', 'scalapack', 'elpa']) 

32def backend_problemsize_kwargs(request): 

33 name = request.param 

34 eigenproblem_size = world.size * 64 

35 if name == 'eigh': 

36 return ScipyDiagonalizer, eigenproblem_size, {'comm': world} 

37 elif name == 'scalapack' or name == 'elpa': 

38 if not compiled_with_sl(): 

39 pytest.skip() 

40 

41 nrows = 2 if world.size > 1 else 1 

42 ncols = world.size // 2 if nrows > 1 else 1 

43 

44 scalapack_kwargs = { 

45 'arraysize': eigenproblem_size, 

46 'grid_nrows': nrows, 

47 'grid_ncols': ncols, 

48 'scalapack_communicator': world, 

49 'blocksize': 32 if world.size == 1 else 64} 

50 

51 if name == 'scalapack': 

52 return ScalapackDiagonalizer, eigenproblem_size, scalapack_kwargs 

53 

54 if name == 'elpa': 

55 if not LibElpa.have_elpa(): 

56 pytest.skip() 

57 return ElpaDiagonalizer, eigenproblem_size, scalapack_kwargs 

58 

59 

60@pytest.mark.parametrize('dtype,', [float, complex]) 

61def test_diagonalizer_eigenproblem_correctness(backend_problemsize_kwargs, 

62 dtype): 

63 is_master_rank = world.rank == 0 

64 ( 

65 diagonalizer_class, 

66 eigenproblem_size, 

67 diagonalizer_kwargs) = backend_problemsize_kwargs 

68 

69 if diagonalizer_class is ScipyDiagonalizer: 

70 diagonalizer = diagonalizer_class(**diagonalizer_kwargs) 

71 elif issubclass(diagonalizer_class, DistributedBlacsDiagonalizer): 

72 diagonalizer = diagonalizer_class(**diagonalizer_kwargs, dtype=dtype) 

73 

74 a, b = prepare_eigensolver_matrices(eigenproblem_size, dtype=dtype) 

75 eps = np.zeros(eigenproblem_size) 

76 a_copy = a.copy() 

77 b_copy = b.copy() 

78 

79 # a_copy contains eigenvectors after this. 

80 diagonalizer.diagonalize( 

81 a_copy, 

82 b_copy, 

83 eps, 

84 debug=False) 

85 

86 if is_master_rank: 

87 assert np.allclose(a @ a_copy, b @ a_copy @ np.diag(eps))